#!/usr/bin/env python3
"""
Display OpenAI model pricing and cost estimates
"""

PRICING = {
    # GPT-5 series
    ("gpt-5.2", "GPT-5.2", "更高质量"): (1.75, 14.00),
    ("gpt-5.1", "GPT-5.1", "高质量"): (1.25, 10.00),
    ("gpt-5", "GPT-5", "高质量"): (1.25, 10.00),
    ("gpt-5-mini", "GPT-5 Mini", "较便宜"): (0.25, 2.00),
    ("gpt-5-nano", "GPT-5 Nano", "最便宜"): (0.05, 0.40),
    # GPT-5 Pro
    ("gpt-5.2-pro", "GPT-5.2 Pro", "顶级质量"): (21.00, 168.00),
    ("gpt-5-pro", "GPT-5 Pro", "顶级质量"): (15.00, 120.00),
    # GPT-4.1 series
    ("gpt-4.1", "GPT-4.1", "平衡选择"): (2.00, 8.00),
    ("gpt-4.1-mini", "GPT-4.1 Mini", "经济实惠"): (0.40, 1.60),
    ("gpt-4.1-nano", "GPT-4.1 Nano", "超低成本"): (0.10, 0.40),
    # GPT-4o series (recommended)
    ("gpt-4o", "GPT-4o ⭐", "质量优秀"): (2.50, 10.00),
    ("gpt-4o-mini", "GPT-4o Mini ⭐⭐", "性价比最高"): (0.15, 0.60),
}

def calculate_cost(input_price, output_price, input_tokens=2000, output_tokens=1500):
    """Calculate cost for given token counts"""
    input_cost = (input_tokens / 1_000_000) * input_price
    output_cost = (output_tokens / 1_000_000) * output_price
    return input_cost + output_cost

def main():
    print("=" * 90)
    print("OPENAI MODEL PRICING & COST ESTIMATES (Dec 2024)")
    print("=" * 90)
    print()
    
    # Assumptions
    input_tokens = 2000
    output_tokens = 1500
    
    print(f"Assumptions: {input_tokens:,} input tokens + {output_tokens:,} output tokens per problem")
    print()
    
    # Header
    print(f"{'Model':<25} {'Input':<12} {'Output':<12} {'Per Prob':<12} {'100 Probs':<12} {'1000 Probs':<12}")
    print("-" * 90)
    
    # Sort by cost per problem
    model_costs = []
    for (model_id, model_name, desc), (input_price, output_price) in PRICING.items():
        cost_per_problem = calculate_cost(input_price, output_price, input_tokens, output_tokens)
        model_costs.append((model_name, input_price, output_price, cost_per_problem, desc))
    
    model_costs.sort(key=lambda x: x[3])  # Sort by cost per problem
    
    for model_name, input_price, output_price, cost_per_problem, desc in model_costs:
        cost_100 = cost_per_problem * 100
        cost_1000 = cost_per_problem * 1000
        
        print(f"{model_name:<25} ${input_price:<11.2f} ${output_price:<11.2f} "
              f"${cost_per_problem:<11.6f} ${cost_100:<11.2f} ${cost_1000:<11.2f}")
    
    print("-" * 90)
    print()
    
    # Recommendations
    print("RECOMMENDATIONS:")
    print("  ⭐⭐ gpt-4o-mini: Best cost-effectiveness for large batches")
    print("  ⭐  gpt-4o: Best quality-to-cost ratio for high-quality datasets")
    print("  💰  gpt-4.1-nano or gpt-5-nano: Cheapest options for experiments")
    print("  🏆  gpt-5-pro/gpt-5.2-pro: Premium options for critical applications")
    print()
    
    # Sample budgets
    print("SAMPLE BUDGETS:")
    print()
    
    recommended = [
        ("gpt-5-nano", "Budget Testing", 0.05, 0.40),
        ("gpt-4.1-nano", "Low Budget", 0.10, 0.40),
        ("gpt-4o-mini", "Recommended ⭐", 0.15, 0.60),
        ("gpt-4.1-mini", "Mid Budget", 0.40, 1.60),
        ("gpt-4o", "High Quality ⭐", 2.50, 10.00),
    ]
    
    budget_amounts = [1.0, 5.0, 10.0, 20.0, 50.0]
    
    for model_id, label, input_p, output_p in recommended:
        cost_per = calculate_cost(input_p, output_p, input_tokens, output_tokens)
        print(f"\n{label} ({model_id}):")
        print(f"  Cost per problem: ${cost_per:.6f}")
        counts = [int(budget / cost_per) for budget in budget_amounts]
        budget_str = " | ".join([f"${b}→{c:,}" for b, c in zip(budget_amounts, counts)])
        print(f"  Problems: {budget_str}")
    
    print()
    print("=" * 90)

if __name__ == "__main__":
    main()