import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# ============================================================
# Problem 2 Complete Analysis - MCM O-Award Standard
# ============================================================

plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['mathtext.fontset'] = 'stix'
plt.rcParams['axes.unicode_minus'] = False

# =========================
# 1. Data Preparation
# =========================
scenarios = ['Gaming', 'Navigation', 'Movie', 'Chatting', 'Screen Off']
power_W = [3.551, 2.954, 2.235, 1.481, 0.517]  # Average power consumption

# Time-to-empty data (hours) for different start SOC
data_matrix = np.array([
    [4.11, 3.05, 2.01, 0.97],   # Gaming
    [5.01, 3.72, 2.45, 1.18],   # Navigation
    [6.63, 4.92, 3.24, 1.56],   # Movie
    [10.02, 7.43, 4.89, 2.36],  # Chatting
    [29.45, 21.85, 14.39, 6.95] # Screen Off
])
start_soc = [100, 75, 50, 25]

# =========================
# 2. Model Prediction vs Observation
# =========================
# Assuming linear discharge model: T = (SOC - z_min) * C / P
# where C = battery capacity (Wh), z_min = 2%

# Estimate battery capacity from Screen Off data (most stable)
# T_screen_off_100 = (100 - 2) * C / 0.517 => C = T * P / 98
C_estimated = 29.45 * 0.517 / 0.98  # ≈ 15.53 Wh

# Predicted time-to-empty
def predict_tte(start_soc, power, capacity=15.53, z_min=2):
    """Predict time-to-empty based on linear model"""
    return (start_soc - z_min) * capacity / power / 100

predicted_matrix = np.zeros_like(data_matrix)
for i, p in enumerate(power_W):
    for j, soc in enumerate(start_soc):
        predicted_matrix[i, j] = predict_tte(soc, p)

# =========================
# 3. Error Analysis & Uncertainty Quantification
# =========================
error_matrix = data_matrix - predicted_matrix
relative_error = error_matrix / data_matrix * 100  # Percentage error

# Calculate RMSE and MAE for each scenario
rmse_per_scenario = np.sqrt(np.mean(error_matrix**2, axis=1))
mae_per_scenario = np.mean(np.abs(error_matrix), axis=1)

print("=" * 60)
print("PROBLEM 2: TIME-TO-EMPTY PREDICTION ANALYSIS")
print("=" * 60)

# =========================
# Table 1: Prediction vs Observation
# =========================
print("\n[Table 1] Time-to-Empty Predictions vs Observations (hours)")
print("-" * 70)
print(f"{'Scenario':<12} | {'SOC=100%':>10} | {'SOC=75%':>10} | {'SOC=50%':>10} | {'SOC=25%':>10}")
print("-" * 70)
for i, s in enumerate(scenarios):
    obs = ' / '.join([f"{data_matrix[i,j]:.2f}" for j in range(4)])
    pred = ' / '.join([f"{predicted_matrix[i,j]:.2f}" for j in range(4)])
    print(f"{s:<12} | Obs:  {data_matrix[i,0]:>5.2f} | {data_matrix[i,1]:>10.2f} | {data_matrix[i,2]:>10.2f} | {data_matrix[i,3]:>10.2f}")
    print(f"{'':12} | Pred: {predicted_matrix[i,0]:>5.2f} | {predicted_matrix[i,1]:>10.2f} | {predicted_matrix[i,2]:>10.2f} | {predicted_matrix[i,3]:>10.2f}")
print("-" * 70)

# =========================
# Table 2: Model Performance Metrics
# =========================
print("\n[Table 2] Model Performance by Scenario")
print("-" * 55)
print(f"{'Scenario':<12} | {'Power(W)':>8} | {'RMSE(h)':>8} | {'MAE(h)':>8} | {'Status':>10}")
print("-" * 55)
for i, s in enumerate(scenarios):
    status = "Good" if rmse_per_scenario[i] < 0.5 else ("Fair" if rmse_per_scenario[i] < 1.0 else "Poor")
    print(f"{s:<12} | {power_W[i]:>8.3f} | {rmse_per_scenario[i]:>8.3f} | {mae_per_scenario[i]:>8.3f} | {status:>10}")
print("-" * 55)

# =========================
# 4. Uncertainty Quantification (95% CI)
# =========================
print("\n[Table 3] Uncertainty Quantification (95% Confidence Interval)")
print("-" * 60)
# Assume 5% measurement uncertainty in power + 3% in capacity
power_uncertainty = 0.05
capacity_uncertainty = 0.03
total_uncertainty = np.sqrt(power_uncertainty**2 + capacity_uncertainty**2)  # ~5.83%

for i, s in enumerate(scenarios):
    t_100 = data_matrix[i, 0]
    ci_lower = t_100 * (1 - 1.96 * total_uncertainty)
    ci_upper = t_100 * (1 + 1.96 * total_uncertainty)
    print(f"{s:<12}: T_100% = {t_100:.2f}h, 95% CI = [{ci_lower:.2f}, {ci_upper:.2f}]h")
print("-" * 60)

# =========================
# 5. Drivers of Rapid Battery Drain
# =========================
print("\n[Analysis 1] Drivers of Rapid Battery Drain")
print("=" * 60)

# Rank by power consumption
sorted_idx = np.argsort(power_W)[::-1]
print("\nRanking by Power Consumption (Greatest to Least Impact):")
print("-" * 50)
for rank, i in enumerate(sorted_idx, 1):
    drain_rate = power_W[i] / C_estimated * 100  # %/hour
    time_factor = data_matrix[4, 0] / data_matrix[i, 0]  # Relative to Screen Off
    print(f"  {rank}. {scenarios[i]:<12}: {power_W[i]:.3f}W ({drain_rate:.1f}%/h)")
    print(f"     → {time_factor:.1f}x faster drain than Screen Off")
print("-" * 50)

# Key drivers identification
print("\nKey Drivers of Rapid Drain:")
print("  • Gaming: GPU rendering + high CPU usage + screen brightness")
print("  • Navigation: GPS module + continuous screen + network")
print("  • Movie: Video decoding + screen backlight + audio")

# =========================
# 6. Activities with Surprisingly Little Impact
# =========================
print("\n[Analysis 2] Activities with Surprisingly Little Model Change")
print("=" * 60)

# Compare model sensitivity
base_time = data_matrix[4, 0]  # Screen Off as baseline
for i, s in enumerate(scenarios):
    reduction = (base_time - data_matrix[i, 0]) / base_time * 100
    expected = (power_W[i] - power_W[4]) / power_W[4] * 100
    surprise = abs(reduction - expected) / expected * 100 if expected > 0 else 0
    
    if i < 4:  # Not Screen Off itself
        if surprise < 20:
            verdict = "As Expected"
        elif reduction < expected * 0.8:
            verdict = "Surprisingly Small Impact"
        else:
            verdict = "Surprisingly Large Impact"
        print(f"{s:<12}: {reduction:>5.1f}% reduction (Expected ~{expected:.0f}% based on power)")
        print(f"             → {verdict}")

print("\n" + "-" * 60)
print("Conclusion on 'Surprisingly Little' Impact:")
print("  • Chatting: Low active screen time → power dominated by idle")
print("  • The OLED display scaling means text-based apps consume")
print("    surprisingly little extra power compared to screen off")
print("  • Background tasks (OS overhead) create a 'floor' effect")
print("-" * 60)

# =========================
# 7. Visualization: Error Analysis Figure
# =========================
fig, axes = plt.subplots(1, 3, figsize=(14, 4.5), dpi=300)

# Nature-style colors
colors = ['#E64B35', '#4DBBD5', '#00A087', '#3C5488', '#F39B7F']

# (a) Prediction vs Observation scatter
ax1 = axes[0]
for i, s in enumerate(scenarios):
    ax1.scatter(data_matrix[i, :], predicted_matrix[i, :], 
                c=colors[i], s=60, label=s, alpha=0.8, edgecolors='black', linewidth=0.5)
max_val = max(data_matrix.max(), predicted_matrix.max())
ax1.plot([0, max_val*1.05], [0, max_val*1.05], 'k--', linewidth=1.5, label='Perfect Fit')
ax1.set_xlabel("Observed Time-to-Empty (h)", fontsize=11, fontweight='bold')
ax1.set_ylabel("Predicted Time-to-Empty (h)", fontsize=11, fontweight='bold')
ax1.legend(fontsize=8, loc='lower right')
ax1.set_xlim(0, max_val*1.05)
ax1.set_ylim(0, max_val*1.05)
ax1.text(0.05, 0.95, '(a)', transform=ax1.transAxes, fontsize=12, fontweight='bold', va='top')
ax1.grid(True, alpha=0.3)

# (b) Relative Error by Scenario
ax2 = axes[1]
x_pos = np.arange(len(scenarios))
mean_rel_error = np.mean(np.abs(relative_error), axis=1)
std_rel_error = np.std(np.abs(relative_error), axis=1)
bars = ax2.bar(x_pos, mean_rel_error, yerr=std_rel_error, capsize=4, 
               color=colors, edgecolor='black', linewidth=1)
ax2.set_xticks(x_pos)
ax2.set_xticklabels(scenarios, rotation=30, ha='right', fontsize=9)
ax2.set_ylabel("Mean Absolute Relative Error (%)", fontsize=11, fontweight='bold')
ax2.axhline(y=10, color='red', linestyle='--', linewidth=1.5, label='10% Threshold')
ax2.legend(fontsize=9)
ax2.text(0.05, 0.95, '(b)', transform=ax2.transAxes, fontsize=12, fontweight='bold', va='top')
ax2.grid(True, alpha=0.3, axis='y')

# (c) Power vs Battery Life (inverse relationship)
ax3 = axes[2]
ax3.scatter(power_W, data_matrix[:, 0], c=colors, s=100, edgecolors='black', linewidth=1, zorder=5)
for i, s in enumerate(scenarios):
    ax3.annotate(s, (power_W[i], data_matrix[i, 0]), 
                 textcoords="offset points", xytext=(5, 5), fontsize=8)

# Fit inverse curve: T = k / P
k_fit = np.mean(np.array(power_W) * data_matrix[:, 0])
p_range = np.linspace(0.3, 4, 100)
t_fit = k_fit / p_range
ax3.plot(p_range, t_fit, 'k--', linewidth=1.5, label=f'$T = {k_fit:.1f}/P$')
ax3.set_xlabel("Power Consumption (W)", fontsize=11, fontweight='bold')
ax3.set_ylabel("Time-to-Empty from 100% (h)", fontsize=11, fontweight='bold')
ax3.legend(fontsize=9)
ax3.text(0.05, 0.95, '(c)', transform=ax3.transAxes, fontsize=12, fontweight='bold', va='top')
ax3.grid(True, alpha=0.3)
ax3.set_xlim(0, 4.5)
ax3.set_ylim(0, 35)

plt.tight_layout()
plt.savefig('problem2_analysis.png', dpi=300, bbox_inches='tight')
plt.savefig('problem2_analysis.pdf', bbox_inches='tight')
print("\n[Figure saved: problem2_analysis.png/pdf]")

print("\n" + "=" * 60)
print("ANALYSIS COMPLETE")
print("=" * 60)