import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import pandas as pd
import os

# ==========================================
# 1. Configuration & Plotting Style Setup
# ==========================================
def configure_plots():
    """Configure Matplotlib to meet academic standards (Times New Roman)."""
    plt.rcParams['font.family'] = 'serif'
    plt.rcParams['font.serif'] = ['Times New Roman']
    plt.rcParams['axes.labelsize'] = 12
    plt.rcParams['xtick.labelsize'] = 10
    plt.rcParams['ytick.labelsize'] = 10
    plt.rcParams['legend.fontsize'] = 10
    plt.rcParams['figure.dpi'] = 150
    plt.rcParams['savefig.dpi'] = 300
    # Ensure output directory exists
    if not os.path.exists('output'):
        os.makedirs('output')

# ==========================================
# 2. Physical Model Class
# ==========================================
class SmartphoneBatteryModel:
    def __init__(self, capacity_mah=4000, temp_c=25, r_int=0.15):
        """
        Initialize the battery model.
        
        Args:
            capacity_mah (float): Design capacity in mAh.
            temp_c (float): Ambient temperature in Celsius.
            r_int (float): Internal resistance in Ohms.
        """
        self.q_design = capacity_mah / 1000.0  # Convert to Ah
        self.temp_k = temp_c + 273.15
        self.r_int = r_int
        
        # Temperature correction for capacity (Arrhenius-like approximation)
        # Reference T = 298.15K (25C). Lower temp -> Lower capacity.
        self.temp_factor = np.exp(0.5 * (1 - 298.15 / self.temp_k))
        # Clamp factor to reasonable bounds (e.g., 0.5 to 1.1)
        self.temp_factor = np.clip(self.temp_factor, 0.1, 1.2)
        
        self.q_eff = self.q_design * self.temp_factor
        
        # Shepherd Model Parameters for OCV (Open Circuit Voltage)
        # V = E0 - K/SOC + A*exp(-B*(1-SOC))
        # Tuned for a typical Li-ion 3.7V/4.2V cell
        self.E0 = 3.4
        self.K = 0.02   # Polarization constant
        self.A = 0.6    # Exponential zone amplitude
        self.B = 20.0   # Exponential zone time constant
        
    def get_ocv(self, soc):
        """
        Calculate Open Circuit Voltage based on SOC.
        Includes a safety clamp for SOC to avoid division by zero.
        """
        soc = np.clip(soc, 0.01, 1.0) # Avoid singularity at SOC=0
        
        # Simplified Shepherd Model + Linear term for better fit
        # V_ocv = Constant + Linear*SOC - Polarization + Exponential_Drop
        v_ocv = 3.2 + 0.6 * soc - (0.05 / soc) + 0.5 * np.exp(-20 * (1 - soc))
        return np.clip(v_ocv, 2.5, 4.3)

    def calculate_current(self, power_watts, soc):
        """
        Calculate current I given Power P and SOC.
        Solves the quadratic equation: P = (V_ocv - I * R_int) * I
        => R_int * I^2 - V_ocv * I + P = 0
        """
        v_ocv = self.get_ocv(soc)
        
        # Quadratic coefficients: a*I^2 + b*I + c = 0
        a = self.r_int
        b = -v_ocv
        c = power_watts
        
        delta = b**2 - 4*a*c
        
        if delta < 0:
            # Power demand exceeds battery capability (voltage collapse)
            # Return a very high current to simulate crash or max out
            return v_ocv / (2 * self.r_int) 
        
        # We want the smaller root (stable operating point)
        # I = (-b - sqrt(delta)) / 2a
        current = (-b - np.sqrt(delta)) / (2 * a)
        return current

    def derivative(self, t, y, power_func):
        """
        The ODE function dy/dt = f(t, y).
        y[0] = SOC (0.0 to 1.0)
        """
        soc = y[0]
        
        if soc <= 0:
            return [0.0] # Battery empty
        
        # Get instantaneous power demand from the scenario function
        p_load = power_func(t)
        
        # Calculate current required to support this power
        i_load = self.calculate_current(p_load, soc)
        
        # d(SOC)/dt = -I / Q_eff
        # Units: I in Amps, Q in Ah, t in Seconds.
        # We need to convert Q to Amp-seconds (Coulombs) -> Q * 3600
        d_soc_dt = -i_load / (self.q_eff * 3600.0)
        
        return [d_soc_dt]

# ==========================================
# 3. Scenario Definitions
# ==========================================
def scenario_idle(t):
    """Scenario A: Idle (Background tasks only). Constant low power."""
    return 0.2  # 200 mW

def scenario_social(t):
    """Scenario B: Social Media (Screen on, fluctuating network)."""
    base = 1.5 # Screen + CPU
    noise = 0.5 * np.sin(t / 60.0) # Fluctuating network usage every minute
    return base + max(0, noise)

def scenario_gaming(t):
    """Scenario C: Heavy Gaming (High CPU/GPU, High Screen)."""
    base = 4.5 # High power
    # Power increases slightly over time due to thermal throttling inefficiency or complexity
    trend = 0.0001 * t 
    return base + trend

# ==========================================
# 4. Simulation & Visualization Logic
# ==========================================
def run_simulation():
    configure_plots()
    
    # Initialize Model
    battery = SmartphoneBatteryModel(capacity_mah=4000, temp_c=25, r_int=0.15)
    
    # Time span: 0 to 24 hours (in seconds)
    t_span = (0, 24 * 3600)
    t_eval = np.linspace(0, 24 * 3600, 1000) # Evaluation points
    
    scenarios = {
        "Idle (Background)": scenario_idle,
        "Social Media": scenario_social,
        "Heavy Gaming": scenario_gaming
    }
    
    results = {}
    
    print(f"{'='*60}")
    print(f"{'Simulation Start':^60}")
    print(f"{'='*60}")
    print(f"Battery Capacity: {battery.q_design*1000:.0f} mAh")
    print(f"Temperature: {battery.temp_k - 273.15:.1f} C")
    print("-" * 60)

    # Solve ODE for each scenario
    for name, p_func in scenarios.items():
        # Solve IVP
        # Stop event: SOC reaches 0
        def battery_empty(t, y): return y[0]
        battery_empty.terminal = True
        
        sol = solve_ivp(
            fun=lambda t, y: battery.derivative(t, y, p_func),
            t_span=t_span,
            y0=[1.0], # Start at 100% SOC
            t_eval=t_eval,
            events=battery_empty,
            method='RK45'
        )
        
        # Post-process to get Voltage and Power for plotting
        soc_vals = sol.y[0]
        time_vals = sol.t
        
        # Re-calculate V and P for visualization
        voltages = []
        powers = []
        for t, s in zip(time_vals, soc_vals):
            p = p_func(t)
            i = battery.calculate_current(p, s)
            v = battery.get_ocv(s) - i * battery.r_int
            voltages.append(v)
            powers.append(p)
            
        results[name] = {
            "time_h": time_vals / 3600.0, # Convert to hours
            "soc": soc_vals * 100.0,      # Convert to %
            "voltage": voltages,
            "power": powers,
            "tte": time_vals[-1] / 3600.0 # Time to empty
        }
        
        print(f"Scenario: {name:<20} | Time-to-Empty: {results[name]['tte']:.2f} hours")

    print(f"{'='*60}")

    # ==========================================
    # 5. Plotting
    # ==========================================
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    
    colors = ['#2ca02c', '#1f77b4', '#d62728'] # Green, Blue, Red
    
    # Plot 1: SOC vs Time
    ax1 = axes[0]
    for (name, data), color in zip(results.items(), colors):
        ax1.plot(data["time_h"], data["soc"], label=f"{name} (TTE={data['tte']:.1f}h)", linewidth=2, color=color)
    ax1.set_title("State of Charge (SOC) vs. Time", fontsize=14, fontweight='bold')
    ax1.set_xlabel("Time (Hours)")
    ax1.set_ylabel("SOC (%)")
    ax1.set_ylim(0, 105)
    ax1.grid(True, linestyle='--', alpha=0.6)
    ax1.legend()

    # Plot 2: Terminal Voltage vs SOC
    ax2 = axes[1]
    for (name, data), color in zip(results.items(), colors):
        # Plot Voltage against SOC (reversed x-axis usually)
        ax2.plot(data["soc"], data["voltage"], label=name, linewidth=2, color=color)
    ax2.set_title("Terminal Voltage vs. SOC", fontsize=14, fontweight='bold')
    ax2.set_xlabel("SOC (%)")
    ax2.set_ylabel("Voltage (V)")
    ax2.invert_xaxis() # Standard battery curve convention
    ax2.grid(True, linestyle='--', alpha=0.6)
    
    # Plot 3: Power Consumption Profile
    ax3 = axes[2]
    for (name, data), color in zip(results.items(), colors):
        ax3.plot(data["time_h"], data["power"], label=name, linewidth=2, color=color, alpha=0.8)
    ax3.set_title("Power Consumption Profile", fontsize=14, fontweight='bold')
    ax3.set_xlabel("Time (Hours)")
    ax3.set_ylabel("Power (Watts)")
    ax3.grid(True, linestyle='--', alpha=0.6)
    
    plt.tight_layout()
    
    # Save and Show
    save_path = os.path.join('output', 'battery_simulation_results.png')
    plt.savefig(save_path)
    print(f"Plot saved to: {save_path}")
    plt.show()

if __name__ == "__main__":
    run_simulation()