import numpy as np
import matplotlib.pyplot as plt

# -------------------------------
# Define the battery model functions
# -------------------------------

def ocv_model(soc, temp):
    """
    Open-circuit voltage model as a function of SOC and temperature.
    For demonstration, we use a simple linear relationship.
    In a real application, this should be replaced with an empirically derived model.
    
    Parameters:
      soc: state-of-charge (0 to 1)
      temp: temperature (°C) [not used in this simple model]
    Returns:
      OCV (V)
    """
    # Example: voltage ranges from 3.0V to 4.2V
    return 3.0 + 1.2 * soc

def battery_model(x, current, dt):
    """
    State propagation model.
    State vector x = [SOC, capacity]
      - SOC: state-of-charge (0 to 1)
      - capacity: battery capacity in Ah
    The SOC decreases with discharge current.
    
    Parameters:
      x: state vector [soc, capacity]
      current: battery current (A) (positive for discharge)
      dt: time interval (seconds)
    Returns:
      Updated state vector.
    """
    soc, capacity = x
    # Note: capacity is in Ah and dt is in seconds, so convert capacity to Coulombs (Ah * 3600)
    soc_new = soc - (current * dt) / (capacity * 3600)
    # For simplicity, assume capacity changes very slowly; we let it remain nearly constant.
    return np.array([soc_new, capacity])

def jacobian_f(x, current, dt):
    """
    Computes the Jacobian of the battery_model function with respect to state x.
    """
    soc, capacity = x
    # f1 = soc - (current*dt)/(capacity*3600)
    df1_dsoc = 1.0
    df1_dcap = (current * dt) / (capacity**2 * 3600)
    # f2 = capacity (assumed constant), so its derivatives are:
    df2_dsoc = 0.0
    df2_dcap = 1.0
    return np.array([[df1_dsoc, df1_dcap],
                     [df2_dsoc, df2_dcap]])

def measurement_model(x, current, temp, R=0.05):
    """
    Measurement model: the terminal voltage is given by the open-circuit voltage
    minus the voltage drop across the internal resistance.
    
    Parameters:
      x: state vector [soc, capacity]
      current: battery current (A)
      temp: temperature (°C)
      R: internal resistance (Ohm)
    Returns:
      Terminal voltage (V)
    """
    soc, _ = x
    return ocv_model(soc, temp) - current * R

def jacobian_h(x, current, temp, R=0.05):
    """
    Computes the Jacobian of the measurement_model with respect to the state x.
    Here we assume the OCV model derivative with respect to SOC is constant (1.2 V per unit SOC).
    """
    # dV/dsoc: derivative of ocv_model with respect to soc (in our simple model, it's constant)
    dV_dsoc = 1.2
    dV_dcap = 0.0  # Voltage is not directly affected by capacity in this simple model.
    return np.array([dV_dsoc, dV_dcap]).reshape(1, -1)

# -------------------------------
# Extended Kalman Filter functions
# ------------------------------- 

def ekf_predict(x, P, current, dt, Q):
    """
    EKF prediction step.
    
    Parameters:
      x: current state estimate
      P: current state covariance matrix
      current: current battery current (A)
      dt: time step (s)
      Q: process noise covariance matrix
    Returns:
      x_pred: predicted state
      P_pred: predicted covariance
    """
    x_pred = battery_model(x, current, dt)
    F = jacobian_f(x, current, dt)
    P_pred = F @ P @ F.T + Q
    return x_pred, P_pred

def ekf_update(x_pred, P_pred, z, current, temp, R_internal=0.05, R_measure=0.01):
    """
    EKF update step.
    
    Parameters:
      x_pred: predicted state
      P_pred: predicted covariance matrix
      z: measured voltage (V)
      current: battery current (A)
      temp: temperature (°C)
      R_internal: internal resistance used in the measurement model
      R_measure: measurement noise variance
    Returns:
      x_updated: updated state estimate
      P_updated: updated covariance matrix
    """
    H = jacobian_h(x_pred, current, temp, R_internal)
    z_pred = measurement_model(x_pred, current, temp, R_internal)
    y = z - z_pred  # innovation or measurement residual
    S = H @ P_pred @ H.T + R_measure  # innovation covariance
    K = P_pred @ H.T @ np.linalg.inv(S)  # Kalman gain
    x_updated = x_pred + (K * y).flatten()  # update state (flatten since K*y is a vector)
    P_updated = (np.eye(len(x_pred)) - K @ H) @ P_pred
    return x_updated, P_updated

# -------------------------------
# Simulation / Filtering Example
# -------------------------------

# Simulation parameters  
dt = 0.25  # time step in seconds
num_steps = 100  # number of simulation steps

# Initialize state:
# Assume starting at 100% SOC and a capacity of 3 Ah (e.g., 3000 mAh)
x_est = np.array([1.0, 3.0])
# Initial state covariance
P_est = np.diag([0.01, 0.01])
# Process noise covariance matrix
Q = np.diag([1e-5, 1e-6])
# Measurement noise variance (for voltage measurement)
R_measure = 0.01
# Internal resistance (Ohm)
R_internal = 0.05

# Generate dummy measurement data
# For example, assume a constant discharge current of 0.5 A and constant temperature (25°C)
current_array = np.full(num_steps, 0.5)
temp_array = np.full(num_steps, 25)
voltage_meas_array = np.zeros(num_steps)
voltage_true_array = np.zeros(num_steps)

# For simulation purposes, we generate a "true" state trajectory.
x_true = np.array([1.0, 3.0])  # initial true state

# Lists to store estimates for plotting
soc_estimates = []
cap_estimates = []

for k in range(num_steps):
    # --- Simulation: generate true state and corresponding measurement ---
    x_true = battery_model(x_true, current_array[k], dt)
    voltage_true = measurement_model(x_true, current_array[k], temp_array[k], R_internal)
    # Simulate measurement noise
    noise = np.random.normal(0, np.sqrt(R_measure))
    voltage_meas = voltage_true + noise
    voltage_true_array[k] = voltage_true 
    voltage_meas_array[k] = voltage_meas

    # --- EKF Predict Step ---
    x_pred, P_pred = ekf_predict(x_est, P_est, current_array[k], dt, Q)

    # --- EKF Update Step ---
    x_est, P_est = ekf_update(x_pred, P_pred, voltage_meas, current_array[k], temp_array[k],
                              R_internal, R_measure)

    soc_estimates.append(x_est[0])
    cap_estimates.append(x_est[1])
    
    # Optional: print the step results
    print(f"Step {k:03d}: Measured Voltage = {voltage_meas:.3f} V, Estimated SOC = {x_est[0]:.4f}, Capacity = {x_est[1]:.4f} Ah")

# -------------------------------
# Plotting results for visualization
# -------------------------------
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(soc_estimates, label='Estimated SOC')
plt.xlabel('Time Step')
plt.ylabel('State of Charge')
plt.title('SOC Estimation')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(cap_estimates, label='Estimated Capacity (Ah)')
plt.xlabel('Time Step')
plt.ylabel('Capacity (Ah)')
plt.title('Capacity (SOH) Estimation')
plt.legend()

plt.tight_layout()
plt.show()