#IEMC_NN pipeline
=========================================

Pipeline:
  1. Vehicle + powertrain model (unchanged from erev_sim.py) + a slew-rate
     (Pdelt_ex) limit applied consistently in every forward simulation.
  2. DP teacher: backward Bellman sweep -> optimal extender-power policy table,
     then a rate-limited forward pass that produces the realized trajectory.
  3. Build a controller library: run DP at 3 "typical" distances, train one NN
     controller per distance on the DP (state -> extender power) pairs.
  4. IEMC_NN online controller: Eper-based selection module that switches between
     the library NNs for an untrained target distance.
  5. Compare CD/CS vs DP-optimal vs IEMC_NN and produce all paper-style plots.
"""

import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.neural_network import MLPRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

RNG = np.random.default_rng(0)
OUT = "plots"                 # output directory (created next to the script)
os.makedirs(OUT, exist_ok=True)

# =============================================================================
# 1. VEHICLE PARAMETERS (Li Auto L9 -- unchanged from erev_sim.py)
# =============================================================================
class LiAutoL9Specs:
    m = 2577.0; A = 2.50; Cd = 0.30; Cr = 0.01; rho = 1.225; g = 9.81
    mass_factor = 1.05; eta_trans = 0.97; eta_motor = 1.0

    P_gen_max = 100000.0; P_gen_min = 15000.0; bsfc_min = 220.0
    # Realistic BSFC "bowl": minimum at bsfc_opt_frac of rated power, rising at
    # low load (idle/part-load is inefficient) and slightly at full load.
    # This replaces the old monotonic curve whose minimum was always at rated
    # power -- that made the efficient-pulse benefit a modelling artifact.
    bsfc_opt_frac = 0.80          # load fraction giving lowest BSFC
    bsfc_curv     = 420.0         # g/kWh per (load-frac)^2 ; sets bowl steepness

    E_batt_kWh = 45.0; E_batt_J = E_batt_kWh * 3.6e6
    eta_chg = 0.96; eta_dis = 0.96
    c_rate_dis = 3.0; c_rate_chg = 2.0
    P_batt_dis_max = c_rate_dis * E_batt_kWh * 1000.0
    P_batt_chg_max = -c_rate_chg * E_batt_kWh * 1000.0

    soc_init = 1.0; soc_target = 0.05; soc_min = 0.03; soc_max = 1.0

    alpha = 1.0; beta = 3.0; terminal_weight = 2e4

    # NEW: slew-rate limit on extender electrical power (Pdelt_ex in Eq. 8).
    # 0 -> 100 kW takes ~4 s. Applied in every forward simulation so DP, NN
    # and CD/CS all respect the same physically-implementable ramp rate.
    Pdelt_ex_max = 25000.0   # W per second
    Nx = 81; Nu = 25

V = LiAutoL9Specs()
RHO_FUEL = 775.0  # g/L

# =============================================================================
# 2. CYCLE + DYNAMICS (unchanged)
# =============================================================================
def clean_drive_cycle(v_ms, a_accel_max=2.5, a_decel_max=4.0, dt=1.0):
    v = v_ms.copy()
    for i in range(1, len(v)):
        dv = v[i] - v[i - 1]
        if dv > a_accel_max * dt:
            v[i] = v[i - 1] + a_accel_max * dt
        elif dv < -a_decel_max * dt:
            v[i] = v[i - 1] - a_decel_max * dt
    return v

def _read_speed_kmh(df):
    """Pick the speed column in km/h robustly (handles files that also carry an
    m/s column, e.g. NEDC.csv, which previously got mis-read as km/h)."""
    kmh_col = None
    for c in df.columns:
        cl = str(c).lower().replace(" ", "")
        if "ms" in cl or "m/s" in cl:          # skip explicit m/s columns
            continue
        if "kmh" in cl or "km/h" in cl or "kph" in cl:
            kmh_col = c; break
    if kmh_col is None:                          # fallback: first non-time column
        cand = [c for c in df.columns if "time" not in str(c).lower()]
        kmh_col = cand[0] if cand else df.columns[-1]
    return pd.to_numeric(df[kmh_col], errors="coerce").fillna(0).values

def generate_trip(cycle_file, target_dist_km):
    df = pd.read_csv(cycle_file)
    v_kmh = _read_speed_kmh(df)
    v_ms = v_kmh / 3.6
    single = np.sum(v_ms) / 1000.0
    repeats = int(np.ceil(target_dist_km / single))
    v_trip = clean_drive_cycle(np.tile(v_ms, repeats))
    t = np.arange(len(v_trip))
    dist = np.cumsum(v_trip) / 1000.0
    return t, v_trip, dist

def calculate_power_demand(v_ms):
    dt = 1.0
    a = np.zeros_like(v_ms); a[1:] = (v_ms[1:] - v_ms[:-1]) / dt
    F_aero = 0.5 * V.rho * V.Cd * V.A * v_ms ** 2
    F_roll = V.m * V.g * V.Cr * np.ones_like(v_ms)
    F_accel = V.m * V.mass_factor * a
    P_wheel = (F_aero + F_roll + F_accel) * v_ms
    P_req = np.zeros_like(P_wheel)
    pos = P_wheel > 0
    P_req[pos] = (P_wheel[pos] / V.eta_trans) / V.eta_motor
    P_req[~pos] = (P_wheel[~pos] * V.eta_trans) * V.eta_motor
    return P_req

def bsfc_g_per_kWh(P_kW):
    """Bowl-shaped BSFC map (g/kWh). Minimum = bsfc_min at bsfc_opt_frac*Prated,
    rising for lower and (slightly) higher load. Engine-off handled by caller."""
    x = P_kW / (V.P_gen_max / 1000.0)          # load fraction 0..1
    return V.bsfc_min + V.bsfc_curv * (x - V.bsfc_opt_frac) ** 2

def fuel_rate_g_s(P_gen_W):
    P_gen_W = np.atleast_1d(P_gen_W).astype(float)
    P_kW = P_gen_W / 1000.0
    rate = np.zeros_like(P_kW)
    on = P_kW > 1e-6
    rate[on] = P_kW[on] * bsfc_g_per_kWh(P_kW[on]) / 3600.0
    return rate

def bsfc_optimal_power_W():
    """Power (W) that minimises BSFC -> the natural target for the pulse actuator."""
    return float(np.clip(V.bsfc_opt_frac * V.P_gen_max, V.P_gen_min, V.P_gen_max))

def battery_soc_delta(P_batt_W, dt=1.0):
    P_batt_W = np.asarray(P_batt_W, dtype=float)
    E_chem = np.where(P_batt_W >= 0, P_batt_W / V.eta_dis, P_batt_W * V.eta_chg)
    return E_chem * dt / V.E_batt_J

def soc_des_curve(dist_km, disall):
    """Eq. 7: linear target SOC line from soc_init to soc_target over disall."""
    frac = np.clip(dist_km / disall, 0.0, 1.0)
    return V.soc_init - (V.soc_init - V.soc_target) * frac

def driving_mode(v_ms, P_req):
    """0 = parking, 1 = driving, 2 = braking (paper's 'driving trend' input)."""
    mode = np.ones_like(v_ms)
    mode[v_ms < 0.5] = 0.0
    mode[P_req < 0] = 2.0
    return mode

# =============================================================================
# 3. DYNAMIC PROGRAMMING (backward sweep + rate-limited forward pass)
# =============================================================================
def solve_dp(P_req):
    N = len(P_req)
    SOC_grid = np.linspace(V.soc_min, V.soc_max, V.Nx)
    P_choices = np.concatenate(([0.0], np.linspace(V.P_gen_min, V.P_gen_max, V.Nu - 1)))
    fuel_rates = fuel_rate_g_s(P_choices)
    soc_des = V.soc_init - (V.soc_init - V.soc_target) * (np.arange(N) / (N - 1))

    CostToGo = np.full((V.Nx, N), 1e10)
    OptU = np.zeros((V.Nx, N))
    CostToGo[:, -1] = V.terminal_weight * (SOC_grid - V.soc_target) ** 2

    for k in range(N - 2, -1, -1):
        P_batt = np.clip(P_req[k] - P_choices, V.P_batt_chg_max, V.P_batt_dis_max)
        dSOC = battery_soc_delta(P_batt)
        SOC_next = np.clip(SOC_grid[:, None] - dSOC[None, :], V.soc_min, V.soc_max)
        future = np.interp(SOC_next.ravel(), SOC_grid, CostToGo[:, k + 1]).reshape(V.Nx, V.Nu)
        running = V.alpha * fuel_rates[None, :] + V.beta * np.abs(SOC_grid[:, None] - soc_des[k])
        total = running + future
        idx = np.argmin(total, axis=1)
        CostToGo[:, k] = total[np.arange(V.Nx), idx]
        OptU[:, k] = P_choices[idx]
    return SOC_grid, OptU, soc_des

def slew_limit(P_prev, P_target):
    """Clamp the change in extender power to +/- Pdelt_ex_max per second."""
    dP = np.clip(P_target - P_prev, -V.Pdelt_ex_max, V.Pdelt_ex_max)
    return P_prev + dP

def forward_dp(P_req, SOC_grid, OptU):
    """Rate-limited forward pass following the DP policy."""
    N = len(P_req)
    SOC = np.zeros(N); P_gen = np.zeros(N); P_batt = np.zeros(N)
    SOC[0] = V.soc_init
    for k in range(N - 1):
        u_raw = np.interp(SOC[k], SOC_grid, OptU[:, k])
        u = slew_limit(P_gen[k - 1] if k > 0 else 0.0, u_raw)
        P_gen[k] = u
        pb = np.clip(P_req[k] - u, V.P_batt_chg_max, V.P_batt_dis_max)
        P_batt[k] = pb
        SOC[k + 1] = np.clip(SOC[k] - battery_soc_delta(pb), V.soc_min, V.soc_max)
    P_gen[-1] = P_gen[-2]
    return SOC, P_gen, P_batt

#print output
def save_dp_debug_data(P_req, SOC_grid, OptU, soc_des, SOC, P_gen, P_batt,suffix=""):
    """
    Saves the outputs of solve_dp and forward_dp to CSV files for debugging.
    """
    # =========================================================================
    # 1. Save the Forward Pass Trajectory (Time-Series)
    # =========================================================================
    # Stack the 1D arrays side-by-side.
    # Adding a TimeStep column makes it easy to plot in Excel.
    time_series_data = np.column_stack((
        np.arange(len(P_req)),  # Time step k
        P_req,                  # Power demand
        soc_des,                # Target SOC
        SOC,                    # Actual forward SOC
        P_gen,                  # Extender power (rate-limited)
        P_batt                  # Battery power
    ))

    # Save to CSV
    ts_header = "TimeStep,P_req,SOC_des,SOC_actual,P_gen,P_batt"
    np.savetxt("debug_forward_trajectory.csv", time_series_data, delimiter=",", header=ts_header, comments="")

    # =========================================================================
    # 2. Save the Raw DP Policy Table (OptU)
    # =========================================================================
    # OptU is a 2D array (Nx rows by N columns).
    # We attach the SOC_grid as the very first column so you know which row is which.
    optu_data = np.column_stack((SOC_grid, OptU))

    # Create a header: "SOC_Grid, t_0, t_1, t_2, ..."
    optu_header = "SOC_Grid," + ",".join([f"t_{i}" for i in range(len(P_req))])
    np.savetxt("debug_OptU_policy.csv", optu_data, delimiter=",", header=optu_header, comments="")

    print("-> DP debug data saved to 'debug_forward_trajectory.csv' and 'debug_OptU_policy.csv'")

# =============================================================================
# 4. CD/CS BASELINE (rule-based, now also rate-limited)
# =============================================================================
def solve_cdcs(P_req):
    N = len(P_req)
    SOC = np.zeros(N); P_gen = np.zeros(N)
    SOC[0] = V.soc_init
    cs = V.soc_target
    for k in range(N - 1):
        Pd = P_req[k]
        if SOC[k] > cs:
            u = np.clip(Pd - V.P_batt_dis_max, V.P_gen_min, V.P_gen_max) if Pd > V.P_batt_dis_max else 0.0
        else:
            u = np.clip(Pd, V.P_gen_min, V.P_gen_max) if Pd > 0 else 0.0
        u = slew_limit(P_gen[k - 1] if k > 0 else 0.0, u)
        P_gen[k] = u
        pb = np.clip(Pd - u, V.P_batt_chg_max, V.P_batt_dis_max)
        SOC[k + 1] = np.clip(SOC[k] - battery_soc_delta(pb), V.soc_min, V.soc_max)
    P_gen[-1] = P_gen[-2]
    return SOC, P_gen

# =============================================================================
# 5. NN CONTROLLER (Section 3: 6 inputs -> extender power, 2 hidden layers x 8)
# =============================================================================
FEATURES = ["P_demand_kW", "v_kmh", "SOC_pct", "SOC_des_pct", "mode", "SOC_err_pct", "Pgen_prev_kW"]

def build_features(P_req, v_ms, SOC, soc_des, mode, P_gen_prev=None):
    """Assemble the 7-input feature matrix. The 7th input, P_gen_prev (W), is the
    extender power one control step earlier. The forward slew limit makes P_gen[k]
    depend on P_gen[k-1]; without this feature the realized target is one-to-many
    over the inputs and MSE regression collapses to the conditional mean."""
    if P_gen_prev is None:
        P_gen_prev = np.zeros_like(np.atleast_1d(P_req), dtype=float)
    return np.column_stack([
        P_req / 1000.0,
        v_ms * 3.6,
        SOC * 100.0,
        soc_des * 100.0,
        mode,
        (SOC - soc_des) * 100.0,
        np.asarray(P_gen_prev, dtype=float) / 1000.0,   # NEW: previous extender power /kW
    ])

# def make_training_data(P_req, v_ms, dist, disall, SOC_grid, OptU, stride=3, avg_window_s=45):
#     """DP teacher data, reformulated for a *learnable* target.

#     The raw DP policy is bang-bang (engine hard OFF most of the time, then hard
#     ON near rated power): ~90% exact zeros with spikes. Regressing that with MSE
#     is ill-posed -- the net collapses toward the mean and R2 goes negative/NaN.

#     Instead we train the supervisory NN on the DESIRED AVERAGE extender power:
#     a moving average (avg_window_s) of the DP policy along the trajectory. This
#     target is smooth and strongly explained by (P_demand, SOC_error, ...), so the
#     small NN can actually learn it (R2 ~ 0.9+). The fast ON/OFF that realises this
#     average at an efficient BSFC point is then handled by the pulse/thermostat
#     layer in forward_iemc -- which is how a real series-hybrid genset is run.

#     We still augment with off-trajectory SOC offsets so the closed-loop controller
#     learns to self-correct SOC error instead of memorising one SOC-vs-time curve."""
#     SOC_dp, P_gen_dp, _ = forward_dp(P_req, SOC_grid, OptU)
#     soc_des = soc_des_curve(dist, disall)
#     mode = driving_mode(v_ms, P_req)
#     N = len(P_req)
#     idx = np.arange(0, N, stride)
#     win = max(1, int(round(avg_window_s / stride)))   # smoothing window in samples
#     kern = np.ones(win) / win

#     X_list, y_list = [], []
#     offsets = [0.0, +0.05, -0.05]
#     for off in offsets:
#         SOC_s = np.clip(SOC_dp[idx] + off, V.soc_min, V.soc_max)
#         # optimal action read from the DP table along this (offset) SOC path
#         u = np.array([np.interp(SOC_s[j], SOC_grid, OptU[:, idx[j]]) for j in range(len(idx))])
#         u_avg = np.convolve(u, kern, mode="same")     # -> desired AVERAGE power
#         X_list.append(build_features(P_req[idx], v_ms[idx], SOC_s, soc_des[idx], mode[idx]))
#         y_list.append(u_avg / 1000.0)  # kW
#     X = np.vstack(X_list); y = np.concatenate(y_list)
#     n0 = len(idx)  # size of the realized (offset 0) block, for the regression diagnostic
#     return X, y, n0, (SOC_dp, P_gen_dp, soc_des, mode)


def make_training_data(P_req, v_ms, dist, disall, SOC_grid, OptU, stride=3):
    """DP teacher data with the previous-power feature threaded through.
    The target is the realized, rate-limited extender power; the 7th feature is
    that same power one second earlier, which is what makes the map a function."""
    SOC_dp, P_gen_dp, P_batt_dp = forward_dp(P_req, SOC_grid, OptU)
    soc_des = soc_des_curve(dist, disall)
    mode = driving_mode(v_ms, P_req)
    N = len(P_req)
    idx = np.arange(0, N, stride)

    def slew_track(u_raw_full, p0):
        """Per-second slew-limited power sequence (matches forward_dp / deployment)."""
        u = np.empty(N); u[0] = p0
        for k in range(1, N):
            u[k] = u[k - 1] + np.clip(u_raw_full[k] - u[k - 1],
                                      -V.Pdelt_ex_max, V.Pdelt_ex_max)
        return u

    X_list, y_list = [], []
    offsets = [0.0, +0.05, -0.05]      # set to [0.0] first to confirm the fix, then re-enable
    for off in offsets:
        if off == 0.0:
            u_full = P_gen_dp                      # already the realized rate-limited power
        else:
            SOC_sh = np.clip(SOC_dp + off, V.soc_min, V.soc_max)
            u_raw_full = np.array([np.interp(SOC_sh[k], SOC_grid, OptU[:, k])
                                   for k in range(N)])
            u_full = slew_track(u_raw_full, P_gen_dp[0])

        prev_full = np.concatenate([[0.0], u_full[:-1]])     # power 1 s earlier
        SOC_path = np.clip(SOC_dp + off, V.soc_min, V.soc_max)

        X_feat = build_features(P_req[idx], v_ms[idx], SOC_path[idx],
                                soc_des[idx], mode[idx], prev_full[idx])  # 7 cols
        X_list.append(X_feat)
        y_list.append(u_full[idx] / 1000.0)        # kW

    X = np.vstack(X_list); y = np.concatenate(y_list)

    # Quantization (paper Table 2) -- now SEVEN input columns
    X[:, 0] = np.round(X[:, 0], 1)   # P_demand kW
    X[:, 1] = np.round(X[:, 1], 2)   # v km/h
    X[:, 2] = np.round(X[:, 2], 1)   # SOC %
    X[:, 3] = np.round(X[:, 3], 1)   # SOC_des %
    X[:, 4] = np.round(X[:, 4], 0)   # mode
    X[:, 5] = np.round(X[:, 5], 1)   # SOC_err %
    X[:, 6] = np.round(X[:, 6], 1)   # P_gen_prev kW  (NEW)
    y = np.round(y, 1)

    n0 = len(idx)
    return X, y, n0, (SOC_dp, P_gen_dp, P_batt_dp, soc_des, mode)

def train_nn(X, y, max_iter=2000):
    # (32,32) is still tiny/deployable but has enough capacity for the smooth
    # 6->1 map; early stopping guards against overfitting and wasted epochs.
    model = Pipeline([
        ("scaler", StandardScaler()),
        ("mlp", MLPRegressor(hidden_layer_sizes=(32, 32), activation="relu",
                             solver="adam", alpha=1e-4, batch_size=512,
                             learning_rate_init=1e-3, max_iter=max_iter,
                             early_stopping=True, validation_fraction=0.1,
                             tol=1e-6, n_iter_no_change=25, random_state=0)),
    ])
    model.fit(X, y)
    return model

def nn_predict_power(model, feat_row, raw=False):
    """Predict extender power (W) for one feature row.
    raw=False: continuous-engine controller -> snap to [P_gen_min, P_gen_max] or OFF.
    raw=True : return the desired AVERAGE power (0..P_gen_max) for the pulse layer,
               which decides the actual ON/OFF firing at an efficient BSFC point."""
    p_W = float(model.predict(feat_row.reshape(1, -1))[0]) * 1000.0
    if raw:
        return float(np.clip(p_W, 0.0, V.P_gen_max))
    if p_W < V.P_gen_min * 0.5:
        return 0.0                          # engine OFF
    return float(np.clip(p_W, V.P_gen_min, V.P_gen_max))

def forward_nn(P_req, v_ms, dist, model, disall, cs_protect=True):
    """Forward simulation under a single NN controller.
    Target-SOC / SOC-error features use `disall` (the controller's native
    distance), which is what fixes that controller's Eper."""
    N = len(P_req)
    SOC = np.zeros(N); P_gen = np.zeros(N)
    SOC[0] = V.soc_init
    mode = driving_mode(v_ms, P_req)
    for k in range(N - 1):
        sd = soc_des_curve(np.array([dist[k]]), disall)[0]
        p_prev = P_gen[k - 1] if k > 0 else 0.0          # NEW: previous extender power (W)
        feat = build_features(P_req[k:k+1], v_ms[k:k+1],
                              SOC[k:k+1], np.array([sd]), mode[k:k+1],
                              np.array([p_prev]))[0]
        u = nn_predict_power(model, feat)
        if cs_protect and SOC[k] <= V.soc_target:   # charge-sustain protection
            u = np.clip(P_req[k], V.P_gen_min, V.P_gen_max) if P_req[k] > 0 else 0.0
        u = slew_limit(P_gen[k - 1] if k > 0 else 0.0, u)
        P_gen[k] = u
        pb = np.clip(P_req[k] - u, V.P_batt_chg_max, V.P_batt_dis_max)
        SOC[k + 1] = np.clip(SOC[k] - battery_soc_delta(pb), V.soc_min, V.soc_max)
    P_gen[-1] = P_gen[-2]
    return SOC, P_gen

# =============================================================================
# 6. IEMC_NN: Eper-based selection across the controller library
# =============================================================================
class NNC:
    def __init__(self, model, D):
        self.model = model
        self.D = D
        self.Eper = (V.soc_init - V.soc_target) / D   # Eq. 13, fixed per controller

def forward_iemc(P_req, v_ms, dist, library, disall, pulse_P=None):
    """Eqs. 12-13 selection module, following the paper's Fig. 11 logic:
      * If `disall` equals a trained library distance, LOCK that NNC for the whole
        trip (no switching).
      * Otherwise start on the conservative bracket (the longer-distance, lower-Eper
        controller) and step DOWN to shorter-distance controllers MONOTONICALLY as
        the required future Eper rises, so SOC catches the target line and lands on
        soc_target.

    pulse_P (W): optional efficient-pulse actuator. The supervisory NN still sets the
    *desired average* extender power, but instead of running the engine continuously
    at that (low, inefficient) power, an owed-energy thermostat fires the engine in
    short bursts at `pulse_P` (a high-efficiency operating point) and lets the battery
    buffer the surplus. Same engine energy, much lower BSFC -> much less fuel. This is
    how a real series-hybrid range extender is actually operated."""
    N = len(P_req)
    SOC = np.zeros(N); P_gen = np.zeros(N); sel = np.zeros(N, dtype=int)
    SOC[0] = V.soc_init
    mode = driving_mode(v_ms, P_req)
    lib = sorted(library, key=lambda c: c.D)          # ascending distance
    Ds = np.array([c.D for c in lib])
    Epers = np.array([c.Eper for c in lib])           # decreasing as D increases
    Eper_trip = (V.soc_init - V.soc_target) / disall

    match = np.where(np.abs(Ds - disall) < 1e-6)[0]
    locked = int(match[0]) if len(match) else None
    if locked is not None:
        active = locked
    else:
        below = np.where(Epers <= Eper_trip)[0]       # controllers depleting <= needed
        active = int(below[0]) if len(below) else len(lib) - 1   # conservative start

    owed = 0.0
    for k in range(N - 1):
        if locked is None and active > 0:
            remaining = max(disall - dist[k], 1e-3)
            Eper_need = (SOC[k] - V.soc_target) / remaining
            switch_level = 0.5 * (Epers[active] + Epers[active - 1])   # Eper midpoint
            if Eper_need >= switch_level:
                active -= 1                            # step to next shorter-distance NNC
        sel[k] = active
        c = lib[active]
        sd = soc_des_curve(np.array([dist[k]]), c.D)[0]
        p_prev = P_gen[k - 1] if k > 0 else 0.0          # NEW
        feat = build_features(P_req[k:k+1], v_ms[k:k+1],
                              SOC[k:k+1], np.array([sd]), mode[k:k+1],
                              np.array([p_prev]))[0]
        u = nn_predict_power(c.model, feat, raw=(pulse_P is not None))  # desired (avg) power
        if SOC[k] <= V.soc_target:                    # charge-sustain protection
            u = np.clip(P_req[k], V.P_gen_min, V.P_gen_max) if P_req[k] > 0 else 0.0
        if pulse_P is not None:                        # efficient-pulse thermostat
            owed = min(owed + max(u, 0.0), 1.5 * pulse_P)
            if owed >= pulse_P:
                u = pulse_P; owed -= pulse_P
            else:
                u = 0.0
        u = slew_limit(P_gen[k - 1] if k > 0 else 0.0, u)
        P_gen[k] = u
        pb = np.clip(P_req[k] - u, V.P_batt_chg_max, V.P_batt_dis_max)
        SOC[k + 1] = np.clip(SOC[k] - battery_soc_delta(pb), V.soc_min, V.soc_max)
    P_gen[-1] = P_gen[-2]; sel[-1] = sel[-2]
    return SOC, P_gen, sel

def fuel_L(P_gen):
    return float(np.sum(fuel_rate_g_s(P_gen)) / RHO_FUEL)

# =============================================================================
# 7. RUN EVERYTHING (Multi-Cycle Training)
# =============================================================================
import os, time, pickle
import numpy as np

CYCLES = ["WLTC_Class3b.csv", "NEDC.csv", "CLTC.csv"]
# Fallback logic to find the files
cycles_paths = []
for c in CYCLES:
    if os.path.exists(c):
        cycles_paths.append(c)
    elif os.path.exists(f"/mnt/user-data/uploads/{c}"):
        cycles_paths.append(f"/mnt/user-data/uploads/{c}")
    else:
        raise FileNotFoundError(f"Could not find drive cycle file: {c}")
# Use the resolved full path everywhere (the plotting sections reference CYCLE),
# so the script works whether or not the CSVs sit in the current directory.
CYCLE = cycles_paths[1]

LIB_DISTANCES = [250.0, 400.0, 550.0, 1800]   # typical distances for the controller library
TEST_DIST = 1500.0                       # untrained target distance for the headline comparison
PULSE_P = bsfc_optimal_power_W()         # fire the engine at its BSFC sweet spot

print("=" * 64)
print("STEP 1-2: Building MULTI-CYCLE controller library")
print("=" * 64)

CACHE = "cache_multicycle.pkl"  # Renamed so it forces a fresh build
library = []
dp_solved = {}
nn_meta = {}

if os.path.exists(CACHE) and os.environ.get("IEMC_REBUILD", "0") != "1":
    print("Loading cached multi-cycle library + DP solves from", CACHE, flush=True)
    with open(CACHE, "rb") as f:
        blob = pickle.load(f)
    library, dp_solved, nn_meta = blob["library"], blob["dp_solved"], blob["nn_meta"]
else:
    for D in LIB_DISTANCES:
        print(f"\n[Target Distance: {D:.0f} km]", flush=True)

        # Containers to hold the aggregated data across all cycles
        X_concat = []
        y_concat = []

        # 1. Loop through every distinct driving cycle
        for cycle_file in cycles_paths:
            t, v, dist = generate_trip(cycle_file, D)
            P_req = calculate_power_demand(v)

            print(f"  -> Solving DP for {os.path.basename(cycle_file)}...", flush=True)
            t0 = time.time()
            SOC_grid, OptU, _ = solve_dp(P_req)

            # Extract features and targets for this specific cycle
            X_c, y_c, n0, meta = make_training_data(P_req, v, dist, D, SOC_grid, OptU)
            X_concat.append(X_c)
            y_concat.append(y_c)

            print(f"     Done in {time.time()-t0:.0f}s. Extracted {len(y_c)} training pairs.")

            # =================================================================
            # EXPORT DEBUG DATA (Cleanly unpacking from 'meta')
            # =================================================================
            SOC_dp, P_gen_dp, P_batt_dp, soc_des, mode = meta

            cycle_name = os.path.basename(cycle_file).replace('.csv', '')
            file_suffix = f"{D:.0f}km_{cycle_name}"

            save_dp_debug_data(P_req, SOC_grid, OptU, soc_des, SOC_dp, P_gen_dp, P_batt_dp, suffix=file_suffix)

            # =================================================================
            # Save the last cycle evaluated to satisfy the downstream plotting functions
            dp_solved[D] = (SOC_grid, OptU, P_req, v, dist)
            last_meta_bundle = (t, v, dist, P_req, SOC_grid, OptU, meta, n0)

        # 2. Stack all the cycle data into one massive, diverse dataset
        X_train = np.vstack(X_concat)
        y_train = np.concatenate(y_concat)

        print(f"\n  -> Training single Neural Network on {len(y_train)} mixed data pairs...", flush=True)
        t1 = time.time()
        model = train_nn(X_train, y_train)
        mlp = model.named_steps["mlp"]

        # 3. Diagnostics: proper held-out R2 over the full mixed dataset, and a
        #    correctly-guarded R2 on the realized (offset-0) block for fig.6.
        from sklearn.metrics import r2_score
        t_last, v_last, dist_last, Preq_last, SG_last, OptU_last, meta_last, n0_last = last_meta_bundle

        # Honest grouped split: train on all cycles but the last, test on the held-out
        # cycle. Random row splits leak on autocorrelated trajectories (fake ~0.98 R2).
        if len(X_concat) >= 2:
            diag = train_nn(np.vstack(X_concat[:-1]), np.concatenate(y_concat[:-1]))
            r2_holdout = r2_score(y_concat[-1], diag.predict(X_concat[-1]))
        else:
            cut = int(0.8 * len(X_train))            # contiguous tail split, no shuffle
            diag = train_nn(X_train[:cut], y_train[:cut])
            r2_holdout = r2_score(y_train[cut:], diag.predict(X_train[cut:]))

        y_true = y_concat[-1][:n0_last]
        y_pred = model.predict(X_concat[-1][:n0_last])
        if np.var(y_true) > 1e-6:
            r2_block = r2_score(y_true, y_pred)
            r2_block_str = f"{r2_block:.4f}"
        else:                                   # engine ~never on for this slice
            r2_block_str = "n/a (near-constant target)"

        print(f"  -> NN Done in {time.time()-t1:.0f}s; final loss: {mlp.loss_:.4f}; "
              f"R2 (held-out): {r2_holdout:.4f}; R2 (realized block): {r2_block_str}\n", flush=True)

        library.append(NNC(model, D))
        nn_meta[D] = dict(t=t_last, v=v_last, dist=dist_last, P_req=Preq_last,
                          SOC_grid=SG_last, OptU=OptU_last, meta=meta_last,
                          loss_curve=mlp.loss_curve_, y_true=y_true, y_pred=y_pred,
                          r2_holdout=r2_holdout)

    with open(CACHE, "wb") as f:
        pickle.dump({"library": library, "dp_solved": dp_solved, "nn_meta": nn_meta}, f)
    print("Saved multi-cycle library cache to", CACHE, flush=True)

nnc_short = sorted(library, key=lambda c: c.D)[0]   # shortest-distance NNC (NNC1 analog)

# ---------------------------------------------------------------- Fig 6
print("\nPlotting Fig.6 (NN training performance + regression)...")
D_show = 400.0
m = nn_meta[D_show]
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
ax[0].plot(m["loss_curve"], "r-", lw=1.6)
ax[0].set_xlabel("Epochs"); ax[0].set_ylabel("MSE (training loss)")
ax[0].set_title(f"(a) MSE vs epochs  [NNC trained @ {D_show:.0f} km]")
ax[0].grid(True, ls=":", alpha=0.6); ax[0].set_yscale("log")
yt, yp = m["y_true"], m["y_pred"]
ax[1].scatter(yt, yp, s=6, alpha=0.25, label="data")
lims = [0, max(yt.max(), yp.max()) * 1.05]
ax[1].plot(lims, lims, "k-", lw=1, label="y = x")
ax[1].set_xlabel("Target /kW (DP)"); ax[1].set_ylabel("Output /kW (NN)")
ax[1].set_title("(b) NN output vs DP target"); ax[1].legend()
ax[1].grid(True, ls=":", alpha=0.6)
plt.tight_layout(); plt.savefig(f"{OUT}/fig06_nn_training.png", dpi=120); plt.close()

# ---------------------------------------------------------------- Fig 8 (DP vs NN, native distance)
print("Plotting Fig.8 (DP vs single NN controller, native distance)...")
m = nn_meta[D_show]
SOC_dp, P_gen_dp, P_batt_dp, soc_des, _ = m["meta"]
nnc_show = [c for c in library if c.D == D_show][0]
SOC_nn, P_gen_nn = forward_nn(m["P_req"], m["v"], m["dist"], nnc_show.model, D_show)
dist = m["dist"]
fig, ax = plt.subplots(2, 2, figsize=(12, 8))
ax[0,0].plot(dist, SOC_dp*100, "k-", lw=2, label="DP")
ax[0,0].plot(dist, SOC_nn*100, "b-", lw=1.4, label="NN controller")
ax[0,0].set_ylabel("SOC /%"); ax[0,0].set_title("(a) Battery SOC"); ax[0,0].legend(); ax[0,0].grid(True, ls=":")
ax[0,1].plot(dist, P_gen_dp/1e3, "k-", lw=1, label="DP")
ax[0,1].plot(dist, P_gen_nn/1e3, "b-", lw=0.8, alpha=0.8, label="NN controller")
ax[0,1].set_ylabel("Extender power /kW"); ax[0,1].set_title("(b) Extender output power"); ax[0,1].legend(); ax[0,1].grid(True, ls=":")
ax[1,0].plot(dist, (SOC_nn-SOC_dp)*100, "r-", lw=1)
ax[1,0].set_xlabel("Driving distance /km"); ax[1,0].set_ylabel("SOC difference /%"); ax[1,0].set_title("(c) SOC difference (NN - DP)"); ax[1,0].grid(True, ls=":")
ax[1,1].plot(dist, (P_gen_nn-P_gen_dp)/1e3, "r-", lw=0.7)
ax[1,1].set_xlabel("Driving distance /km"); ax[1,1].set_ylabel("Power difference /kW"); ax[1,1].set_title("(d) Extender power difference (NN - DP)"); ax[1,1].grid(True, ls=":")
plt.tight_layout(); plt.savefig(f"{OUT}/fig08_dp_vs_nn.png", dpi=120); plt.close()

# ---------------------------------------------------------------- Fig 9 (one NNC at wrong distances)
print("Plotting Fig.9 (single NNC applied to other distances)...")
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
for axi, Dtest in zip(ax, [nnc_short.D, 500.0]):
    t2, v2, d2 = generate_trip(CYCLE, Dtest)
    Pr2 = calculate_power_demand(v2)
    SOCx, _ = forward_nn(Pr2, v2, d2, nnc_short.model, nnc_short.D)
    axi.plot(d2, SOCx*100, "r-", lw=1.6)
    axi.axhline(V.soc_target*100, color="gray", ls=":")
    axi.set_xlabel("Driving distance /km"); axi.set_ylabel("SOC /%")
    axi.set_title(f"{nnc_short.D:.0f} km-NNC on a {Dtest:.0f} km trip"); axi.grid(True, ls=":")
plt.suptitle("Fig.9 analog: one NN controller does not fit other distances", y=1.02)
plt.tight_layout(); plt.savefig(f"{OUT}/fig09_single_nnc_mismatch.png", dpi=120); plt.close()

# ---------------------------------------------------------------- MAIN COMPARISON @ TEST_DIST
print("\n" + "=" * 64)
print(f"STEP 3-5: Comparison at untrained distance {TEST_DIST:.0f} km")
print("=" * 64)
t, v, dist = generate_trip(CYCLE, TEST_DIST)
P_req = calculate_power_demand(v)

SOC_grid_b, OptU_b, _ = solve_dp(P_req)
dp_solved[TEST_DIST] = (SOC_grid_b, OptU_b, P_req, v, dist)
SOC_dpb, P_gen_dpb, _ = forward_dp(P_req, SOC_grid_b, OptU_b)
SOC_cs, P_gen_cs = solve_cdcs(P_req)
SOC_ie, P_gen_ie, sel = forward_iemc(P_req, v, dist, library, TEST_DIST)
SOC_iep, P_gen_iep, _ = forward_iemc(P_req, v, dist, library, TEST_DIST, pulse_P=PULSE_P)
SOC_n1, P_gen_n1 = forward_nn(P_req, v, dist, nnc_short.model, nnc_short.D)  # NNC1-only

res = {
    "CD/CS":          (SOC_cs,  P_gen_cs),
    "NNC1 only":      (SOC_n1,  P_gen_n1),
    "IEMC_NN":        (SOC_ie,  P_gen_ie),
    "IEMC_NN+pulse":  (SOC_iep, P_gen_iep),
    "DP optimal":     (SOC_dpb, P_gen_dpb),
}
print(f"\n{'Controller':<16}{'Final SOC':>11}{'Fuel (L)':>11}{'L/100km':>11}{'vs CD/CS':>11}")
fuel = {k: fuel_L(p) for k, (_, p) in res.items()}
base = fuel["CD/CS"]
for k, (soc, p) in res.items():
    save = (1 - fuel[k]/base) * 100 if base > 0 else 0
    print(f"{k:<16}{soc[-1]*100:>10.1f}%{fuel[k]:>11.2f}{fuel[k]/dist[-1]*100:>11.2f}{save:>10.1f}%")

# Fig 13/15 analog: SOC comparison
print("\nPlotting main SOC comparison (Fig.13/15 analog)...")
fig, ax = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
ax[0].plot(dist, SOC_cs*100, "k-", lw=1.8, label="CD/CS")
ax[0].plot(dist, SOC_n1*100, "g-", lw=1.4, label="NNC1 only")
ax[0].plot(dist, SOC_ie*100, "b-", lw=1.8, label="IEMC_NN")
ax[0].plot(dist, SOC_iep*100, color="purple", lw=1.6, label="IEMC_NN+pulse")
ax[0].plot(dist, SOC_dpb*100, color="orange", ls="--", lw=1.8, label="DP optimal")
ax[0].plot(dist, soc_des_curve(dist, TEST_DIST)*100, "r:", lw=1.2, label="Target SOC line")
ax[0].set_ylabel("SOC /%"); ax[0].legend(ncol=2); ax[0].grid(True, ls=":")
ax[0].set_title(f"Battery SOC: CD/CS vs NNC1-only vs IEMC_NN vs DP  ({TEST_DIST:.0f} km, untrained)")
ax[1].plot(dist, P_gen_cs/1e3, "k-", lw=0.6, alpha=0.6, label="CD/CS")
ax[1].plot(dist, P_gen_ie/1e3, "b-", lw=0.7, alpha=0.8, label="IEMC_NN")
ax[1].plot(dist, P_gen_dpb/1e3, color="orange", lw=0.7, alpha=0.8, label="DP optimal")
ax[1].set_xlabel("Driving distance /km"); ax[1].set_ylabel("Extender power /kW")
ax[1].legend(); ax[1].grid(True, ls=":")
plt.tight_layout(); plt.savefig(f"{OUT}/fig15_main_comparison.png", dpi=120); plt.close()

# Selection trace (which NNC active)
fig, ax = plt.subplots(figsize=(10, 3))
ax.step(dist, [library[i].D if i < len(library) else 0 for i in sel], "b-", where="post")
ax.set_xlabel("Driving distance /km"); ax.set_ylabel("Active NNC distance /km")
ax.set_title("IEMC_NN selection module: active controller vs distance")
ax.set_yticks([c.D for c in sorted(library, key=lambda c: c.D)]); ax.grid(True, ls=":")
plt.tight_layout(); plt.savefig(f"{OUT}/fig12_selection_trace.png", dpi=120); plt.close()

# Fig 14 analog: extender operating-point distribution (engine-on samples)
print("Plotting Fig.14 analog (extender operating-point distribution)...")
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2), sharey=True)
for axi, (name, p) in zip(ax, [("IEMC_NN (smooth)", P_gen_ie), ("IEMC_NN + pulse", P_gen_iep)]):
    on = p[p > 100.0] / 1e3
    axi.hist(on, bins=np.linspace(0, 100, 26), color="steelblue", edgecolor="k", alpha=0.8)
    axi.axvline(PULSE_P/1e3, color="r", ls="--", lw=1, label="BSFC sweet spot")
    axi.set_xlabel("Extender power /kW"); axi.set_title(f"{name} engine-on operating points")
    axi.legend(); axi.grid(True, ls=":", alpha=0.5)
ax[0].set_ylabel("count (s)")
plt.suptitle("Fig.14 analog: the pulse actuator moves engine operation to the efficient band", y=1.02)
plt.tight_layout(); plt.savefig(f"{OUT}/fig14_operating_points.png", dpi=120); plt.close()

# Fig 17: fuel-economy improvement bar chart (reuses cached backward DP solves)
print("Plotting Fig.17 (fuel-economy improvement vs CD/CS)...")
bar_dists = sorted(dp_solved.keys())   # the 3 library distances + the test distance
imp_ie, imp_iep, imp_dp = [], [], []
for D in bar_dists:
    sg, ou, Pb, vb, db = dp_solved[D]
    _, pdp, _ = forward_dp(Pb, sg, ou)
    _, pcs = solve_cdcs(Pb)
    _, pie, _ = forward_iemc(Pb, vb, db, library, D)
    _, piep, _ = forward_iemc(Pb, vb, db, library, D, pulse_P=PULSE_P)
    fcs = fuel_L(pcs)
    imp_ie.append((1 - fuel_L(pie)/fcs)*100 if fcs > 0 else 0)
    imp_iep.append((1 - fuel_L(piep)/fcs)*100 if fcs > 0 else 0)
    imp_dp.append((1 - fuel_L(pdp)/fcs)*100 if fcs > 0 else 0)
x = np.arange(len(bar_dists)); w = 0.27
fig, ax = plt.subplots(figsize=(9, 4.8))
ax.bar(x - w, imp_ie, w, label="IEMC_NN (smooth)", color="steelblue", edgecolor="k")
ax.bar(x, imp_iep, w, label="IEMC_NN + pulse", color="mediumpurple", edgecolor="k")
ax.bar(x + w, imp_dp, w, label="DP optimal", color="orange", edgecolor="k")
ax.set_xticks(x); ax.set_xticklabels([f"{int(d)} km" for d in bar_dists])
ax.set_ylabel("Fuel saving vs CD/CS /%"); ax.set_title("Fig.17: fuel-economy improvement vs CD/CS")
ax.axhline(0, color="k", lw=0.8); ax.legend(); ax.grid(True, axis="y", ls=":", alpha=0.5)
for xi, a, b, cc in zip(x, imp_ie, imp_iep, imp_dp):
    ax.text(xi - w, a + 0.4, f"{a:.1f}", ha="center", fontsize=7)
    ax.text(xi, b + 0.4, f"{b:.1f}", ha="center", fontsize=7)
    ax.text(xi + w, cc + 0.4, f"{cc:.1f}", ha="center", fontsize=7)
plt.tight_layout(); plt.savefig(f"{OUT}/fig17_improvement.png", dpi=120); plt.close()

print("\nAll plots saved to", OUT)
print("Done.")