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

# ==========================================
# 1. Configuration
# ==========================================
def configure_plots():
    plt.rcParams['font.family'] = 'serif'
    plt.rcParams['font.serif'] = ['Times New Roman']
    plt.rcParams['axes.unicode_minus'] = False
    plt.rcParams['font.size'] = 12
    plt.rcParams['figure.dpi'] = 150

# ==========================================
# 2. Simplified Battery Model for Sensitivity
# ==========================================
class FastBatteryModel:
    def __init__(self, capacity_mah=4000, temp_c=25, r_int=0.15, signal_dbm=-90):
        self.q_design = capacity_mah / 1000.0
        self.temp_k = temp_c + 273.15
        self.r_int = r_int
        self.signal = signal_dbm
        
        # Temp correction
        self.temp_factor = np.clip(np.exp(0.6 * (1 - 298.15 / self.temp_k)), 0.1, 1.2)
        self.q_eff = self.q_design * self.temp_factor

    def estimate_tte(self, load_power_watts):
        """
        Estimate TTE using average current approximation to save time complexity.
        TTE ~ Q_eff / I_avg
        Where I_avg is solved from P = V_avg * I - I^2 * R
        """
        # Signal power penalty (simplified exponential model)
        # Baseline -90dBm. If -110dBm, power increases significantly.
        sig_penalty = 0.0
        if self.signal < -90:
            sig_penalty = 0.5 * ((-90 - self.signal) / 20.0)**2
            
        total_power = load_power_watts + sig_penalty
        
        # Average Voltage approximation (3.7V nominal)
        # We solve: Total_Power = (V_nom - I * R_int) * I
        # R * I^2 - V_nom * I + P = 0
        v_nom = 3.7
        a = self.r_int
        b = -v_nom
        c = total_power
        
        delta = b**2 - 4*a*c
        if delta < 0:
            return 0.0 # Voltage collapse, immediate shutdown
            
        i_avg = (-b - np.sqrt(delta)) / (2*a)
        
        # TTE in hours
        tte = self.q_eff / i_avg
        return tte

# ==========================================
# 3. Sensitivity Analysis Logic (OAT)
# ==========================================
def run_sensitivity_analysis():
    configure_plots()
    
    # Baseline Parameters
    base_params = {
        'Load Power (W)': 1.5,
        'Temperature (°C)': 25.0,
        'Internal R (Ω)': 0.15,
        'Signal (dBm)': -90.0
    }
    
    # Perturbation range (+/- 20%)
    # Note: For Signal and Temp, we use additive perturbation for physical meaning
    perturbations = [-0.2, 0.2] 
    
    results = []
    
    # 1. Calculate Baseline TTE
    base_model = FastBatteryModel(
        temp_c=base_params['Temperature (°C)'],
        r_int=base_params['Internal R (Ω)'],
        signal_dbm=base_params['Signal (dBm)']
    )
    base_tte = base_model.estimate_tte(base_params['Load Power (W)'])
    
    print(f"Baseline TTE: {base_tte:.4f} hours")
    
    # 2. Iterate parameters
    for param_name, base_val in base_params.items():
        row = {'Parameter': param_name}
        
        for p in perturbations:
            # Calculate new parameter value
            if param_name == 'Temperature (°C)':
                # For temp, +/- 20% of Celsius is weird, let's do +/- 10 degrees
                new_val = base_val + (10 if p > 0 else -10)
                val_label = f"{new_val}°C"
            elif param_name == 'Signal (dBm)':
                # For signal, +/- 20% dBm is weird, let's do +/- 20 dBm
                new_val = base_val + (20 if p > 0 else -20)
                val_label = f"{new_val}dBm"
            else:
                # Standard percentage
                new_val = base_val * (1 + p)
                val_label = f"{new_val:.2f}"
            
            # Construct model with new param
            # (Copy base params first)
            current_params = base_params.copy()
            current_params[param_name] = new_val
            
            model = FastBatteryModel(
                temp_c=current_params['Temperature (°C)'],
                r_int=current_params['Internal R (Ω)'],
                signal_dbm=current_params['Signal (dBm)']
            )
            
            new_tte = model.estimate_tte(current_params['Load Power (W)'])
            
            # Calculate % change in TTE
            pct_change = (new_tte - base_tte) / base_tte * 100
            
            if p < 0:
                row['Low_Change_%'] = pct_change
                row['Low_Val'] = val_label
            else:
                row['High_Change_%'] = pct_change
                row['High_Val'] = val_label
                
        results.append(row)
        
    df = pd.DataFrame(results)
    
    # ==========================================
    # 4. Visualization (Tornado Plot)
    # ==========================================
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Create bars
    y_pos = np.arange(len(df))
    
    # High perturbation bars
    rects1 = ax.barh(y_pos, df['High_Change_%'], align='center', height=0.4, color='#d62728', label='High Perturbation')
    # Low perturbation bars
    rects2 = ax.barh(y_pos, df['Low_Change_%'], align='center', height=0.4, color='#1f77b4', label='Low Perturbation')
    
    # Styling
    ax.set_yticks(y_pos)
    ax.set_yticklabels(df['Parameter'])
    ax.invert_yaxis()  # Labels read top-to-bottom
    ax.set_xlabel('Change in Time-to-Empty (TTE) [%]')
    ax.set_title('Sensitivity Analysis: Tornado Diagram (Impact on Battery Life)', fontweight='bold')
    ax.axvline(0, color='black', linewidth=0.8, linestyle='--')
    ax.grid(True, axis='x', linestyle='--', alpha=0.5)
    ax.legend()
    
    # Add value labels
    def autolabel(rects, is_left=False):
        for rect in rects:
            width = rect.get_width()
            label_x = width + (1 if width > 0 else -1) * 0.5
            ha = 'left' if width > 0 else 'right'
            ax.text(label_x, rect.get_y() + rect.get_height()/2, 
                    f'{width:.1f}%', ha=ha, va='center', fontsize=9)

    autolabel(rects1)
    autolabel(rects2)
    
    plt.tight_layout()
    plt.savefig('sensitivity_tornado.png')
    plt.show()
    
    print("\nSensitivity Analysis Complete.")
    print(df[['Parameter', 'Low_Change_%', 'High_Change_%']])

if __name__ == "__main__":
    run_sensitivity_analysis()