analyze_actions_phase3.py

"""
Action Analysis for Phase 3 (Vectorized) Model
Compare with Phase 2.3 (PER) performance
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).parent / 'src'))

from fleet_environment_v05 import MultiBridgeFleetV05, FleetConfig

def analyze_model_performance(model_dir, n_episodes=30):
    """Analyze trained model performance over multiple episodes"""
    
    # Load models (move to CPU for inference)
    urban_agent = torch.load(Path(model_dir) / "models" / "urban_agent_final.pt", weights_only=False, map_location='cpu')
    rural_agent = torch.load(Path(model_dir) / "models" / "rural_agent_final.pt", weights_only=False, map_location='cpu')
    urban_agent.eval()
    rural_agent.eval()
    
    # Create environment
    cfg = FleetConfig()
    env = MultiBridgeFleetV05(cfg)
    
    # Statistics collection
    total_rewards = []
    urban_rewards = []
    rural_rewards = []
    total_costs = []
    urban_budget_usage = []
    rural_budget_usage = []
    
    urban_action_counts = np.zeros(5)  # 5 actions per bridge
    rural_action_counts = np.zeros(8)  # 8 strategies
    
    print(f"\nAnalyzing {model_dir} over {n_episodes} episodes...")
    print("="*80)
    
    for ep in range(n_episodes):
        urban_state, rural_state = env.reset()
        episode_reward = 0
        episode_cost = 0
        urban_spent = 0
        rural_spent = 0
        
        for year in range(30):  # 30-year simulation
            # Urban actions
            with torch.no_grad():
                urban_state_t = torch.FloatTensor(urban_state).unsqueeze(0)
                urban_q_values = urban_agent(urban_state_t)
                urban_actions = urban_q_values.argmax(dim=2).squeeze().cpu().numpy()
            
            # Rural action
            with torch.no_grad():
                rural_state_t = torch.FloatTensor(rural_state).unsqueeze(0)
                rural_q_values = rural_agent(rural_state_t)
                rural_action = rural_q_values.argmax(dim=1).item()
            
            # Count actions
            for action in urban_actions:
                urban_action_counts[action] += 1
            rural_action_counts[rural_action] += 1
            
            # Step environment
            (next_urban_state, next_rural_state), reward, done, info = env.step(
                urban_actions, rural_action
            )
            
            episode_reward += reward
            episode_cost += info['total_cost']
            urban_spent += info.get('urban_spent', 0)
            rural_spent += info.get('rural_spent', 0)
            
            urban_state = next_urban_state
            rural_state = next_rural_state
            
            if done:
                break
        
        total_rewards.append(episode_reward)
        urban_rewards.append(info.get('urban_reward', 0))
        rural_rewards.append(info.get('rural_reward', 0))
        total_costs.append(episode_cost)
        
        # Calculate budget usage
        urban_budget = (cfg.total_annual_budget * cfg.urban_budget_share) * 30
        rural_budget = (cfg.total_annual_budget * cfg.rural_budget_share) * 30
        urban_budget_usage.append((urban_spent / urban_budget) * 100)
        rural_budget_usage.append((rural_spent / rural_budget) * 100)
    
    # Normalize action counts
    urban_action_counts = (urban_action_counts / urban_action_counts.sum()) * 100
    rural_action_counts = (rural_action_counts / rural_action_counts.sum()) * 100
    
    results = {
        'total_rewards': total_rewards,
        'urban_rewards': urban_rewards,
        'rural_rewards': rural_rewards,
        'total_costs': total_costs,
        'urban_budget_usage': urban_budget_usage,
        'rural_budget_usage': rural_budget_usage,
        'urban_action_counts': urban_action_counts,
        'rural_action_counts': rural_action_counts
    }
    
    return results

def print_comparison(phase2_results, phase3_results):
    """Print detailed comparison between Phase 2.3 and Phase 3"""
    
    print("\n" + "="*80)
    print("PHASE 2.3 (PER) vs PHASE 3 (VECTORIZED) COMPARISON")
    print("="*80)
    
    print("\n--- 30-Episode Performance (30-year simulations) ---")
    print(f"{'Metric':<30} {'Phase 2.3':>15} {'Phase 3':>15} {'Difference':>15}")
    print("-"*80)
    
    # Total Reward
    p2_reward = np.mean(phase2_results['total_rewards'])
    p3_reward = np.mean(phase3_results['total_rewards'])
    print(f"{'Cumulative Reward:':<30} {p2_reward:>15.2f} {p3_reward:>15.2f} {p3_reward-p2_reward:>+15.2f}")
    
    # Urban Reward
    p2_urban = np.mean(phase2_results['urban_rewards'])
    p3_urban = np.mean(phase3_results['urban_rewards'])
    print(f"{'Urban Reward:':<30} {p2_urban:>15.2f} {p3_urban:>15.2f} {p3_urban-p2_urban:>+15.2f}")
    
    # Rural Reward
    p2_rural = np.mean(phase2_results['rural_rewards'])
    p3_rural = np.mean(phase3_results['rural_rewards'])
    print(f"{'Rural Reward:':<30} {p2_rural:>15.2f} {p3_rural:>15.2f} {p3_rural-p2_rural:>+15.2f}")
    
    # Total Cost
    p2_cost = np.mean(phase2_results['total_costs'])
    p3_cost = np.mean(phase3_results['total_costs'])
    print(f"{'Total Cost:':<30} {p2_cost:>15.2f} {p3_cost:>15.2f} {p3_cost-p2_cost:>+15.2f}")
    
    # Budget Usage
    p2_urban_budget = np.mean(phase2_results['urban_budget_usage'])
    p3_urban_budget = np.mean(phase3_results['urban_budget_usage'])
    print(f"{'Urban Budget Usage (%):':<30} {p2_urban_budget:>15.2f} {p3_urban_budget:>15.2f} {p3_urban_budget-p2_urban_budget:>+15.2f}")
    
    p2_rural_budget = np.mean(phase2_results['rural_budget_usage'])
    p3_rural_budget = np.mean(phase3_results['rural_budget_usage'])
    print(f"{'Rural Budget Usage (%):':<30} {p2_rural_budget:>15.2f} {p3_rural_budget:>15.2f} {p3_rural_budget-p2_rural_budget:>+15.2f}")
    
    print("\n--- Action Distribution ---")
    print("\nUrban Actions (% of total):")
    action_names = ['Do Nothing', 'Routine', 'Preventive', 'Essential', 'Rebuild']
    for i, name in enumerate(action_names):
        p2_pct = phase2_results['urban_action_counts'][i]
        p3_pct = phase3_results['urban_action_counts'][i]
        print(f"  {name:<15} Phase 2.3: {p2_pct:>6.2f}%  Phase 3: {p3_pct:>6.2f}%  Diff: {p3_pct-p2_pct:>+6.2f}%")
    
    print("\nRural Strategies (% of total):")
    strategy_names = ['Do Nothing', 'Light', 'Medium', 'Heavy', 'Very Heavy', 
                      'Critical', 'Emergency', 'Full Rebuild']
    for i, name in enumerate(strategy_names):
        p2_pct = phase2_results['rural_action_counts'][i]
        p3_pct = phase3_results['rural_action_counts'][i]
        print(f"  {name:<15} Phase 2.3: {p2_pct:>6.2f}%  Phase 3: {p3_pct:>6.2f}%  Diff: {p3_pct-p2_pct:>+6.2f}%")
    
    print("\n" + "="*80)

def plot_comparison(phase2_results, phase3_results, output_dir):
    """Create comparison visualizations"""
    
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    
    # Total Rewards
    ax = axes[0, 0]
    data = [phase2_results['total_rewards'], phase3_results['total_rewards']]
    bp = ax.boxplot(data, labels=['Phase 2.3\n(PER)', 'Phase 3\n(Vectorized)'], 
                     patch_artist=True)
    for patch, color in zip(bp['boxes'], ['lightblue', 'lightgreen']):
        patch.set_facecolor(color)
    ax.set_ylabel('Cumulative Reward')
    ax.set_title('Total Reward Comparison')
    ax.grid(True, alpha=0.3, axis='y')
    
    # Urban Rewards
    ax = axes[0, 1]
    data = [phase2_results['urban_rewards'], phase3_results['urban_rewards']]
    bp = ax.boxplot(data, labels=['Phase 2.3', 'Phase 3'], patch_artist=True)
    for patch, color in zip(bp['boxes'], ['lightblue', 'lightgreen']):
        patch.set_facecolor(color)
    ax.set_ylabel('Urban Reward')
    ax.set_title('Urban Performance')
    ax.grid(True, alpha=0.3, axis='y')
    
    # Rural Rewards
    ax = axes[0, 2]
    data = [phase2_results['rural_rewards'], phase3_results['rural_rewards']]
    bp = ax.boxplot(data, labels=['Phase 2.3', 'Phase 3'], patch_artist=True)
    for patch, color in zip(bp['boxes'], ['lightblue', 'lightgreen']):
        patch.set_facecolor(color)
    ax.set_ylabel('Rural Reward')
    ax.set_title('Rural Performance')
    ax.grid(True, alpha=0.3, axis='y')
    
    # Budget Usage
    ax = axes[1, 0]
    x = np.arange(2)
    width = 0.35
    urban_usage = [np.mean(phase2_results['urban_budget_usage']), 
                   np.mean(phase3_results['urban_budget_usage'])]
    rural_usage = [np.mean(phase2_results['rural_budget_usage']), 
                   np.mean(phase3_results['rural_budget_usage'])]
    ax.bar(x - width/2, urban_usage, width, label='Urban', alpha=0.7)
    ax.bar(x + width/2, rural_usage, width, label='Rural', alpha=0.7)
    ax.set_ylabel('Budget Usage (%)')
    ax.set_title('Budget Utilization')
    ax.set_xticks(x)
    ax.set_xticklabels(['Phase 2.3', 'Phase 3'])
    ax.legend()
    ax.grid(True, alpha=0.3, axis='y')
    
    # Urban Action Distribution
    ax = axes[1, 1]
    action_names = ['Do\nNothing', 'Routine', 'Preventive', 'Essential', 'Rebuild']
    x = np.arange(len(action_names))
    width = 0.35
    ax.bar(x - width/2, phase2_results['urban_action_counts'], width, 
           label='Phase 2.3', alpha=0.7)
    ax.bar(x + width/2, phase3_results['urban_action_counts'], width, 
           label='Phase 3', alpha=0.7)
    ax.set_ylabel('Percentage (%)')
    ax.set_title('Urban Action Distribution')
    ax.set_xticks(x)
    ax.set_xticklabels(action_names, fontsize=9)
    ax.legend()
    ax.grid(True, alpha=0.3, axis='y')
    
    # Rural Strategy Distribution
    ax = axes[1, 2]
    strategy_names = ['Do\nNothing', 'Light', 'Medium', 'Heavy', 'Very\nHeavy', 
                      'Critical', 'Emergency', 'Full\nRebuild']
    x = np.arange(len(strategy_names))
    ax.bar(x - width/2, phase2_results['rural_action_counts'], width, 
           label='Phase 2.3', alpha=0.7)
    ax.bar(x + width/2, phase3_results['rural_action_counts'], width, 
           label='Phase 3', alpha=0.7)
    ax.set_ylabel('Percentage (%)')
    ax.set_title('Rural Strategy Distribution')
    ax.set_xticks(x)
    ax.set_xticklabels(strategy_names, fontsize=8)
    ax.legend()
    ax.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    output_path = Path(output_dir)
    plt.savefig(output_path / 'phase_comparison_actions.png', dpi=300, bbox_inches='tight')
    print(f"\nSaved comparison plot to {output_path / 'phase_comparison_actions.png'}")
    plt.close()

def main():
    # Analyze Phase 2.3 model
    phase2_dir = 'long_train_phase2_per'
    phase2_results = analyze_model_performance(phase2_dir, n_episodes=30)
    
    # Analyze Phase 3 model
    phase3_dir = 'long_train_phase3_vectorized'
    phase3_results = analyze_model_performance(phase3_dir, n_episodes=30)
    
    # Print comparison
    print_comparison(phase2_results, phase3_results)
    
    # Plot comparison
    plot_comparison(phase2_results, phase3_results, phase3_dir)
    
    print("\n✓ Analysis complete!")

if __name__ == '__main__':
    main()