%% TD - Battery Prognostics with Deep LSTM (Sliding Window, Multi-Battery)
% Made by: Dr. Mayank S JHA , 3rd Feb, 2026.
% Goal ---- Train a deep LSTM network to forecast battery capacity using a sliding window approach:
%
% Why "multi-battery" training? ----------------------------- In realistic battery prognostics, we do NOT train only on the early life of one battery and expect accurate end-of-life forecasts. Instead, we train on many batteries (full life / near full life), then test on a NEW battery (a new trajectory).
% This script: 1) Generates synthetic degradation data for MANY batteries.
% 2) Builds a supervised learning dataset using sliding windows.
% 3) Splits the data BY BATTERY ID (train/val/test batteries are different).
% 4) Trains a deep LSTM (stacked LSTMs + FC + dropout) similar to MATLAB docs.
%
%  5) Evaluates L-step-ahead forecasts and plots results.
%

clear; clc; close all;
rng(0); % reproducibility for teaching/demo

%
%% =========================
% PART A) Intuition Plots ========================== We first visualize how capacity decreases with cycle number and how temperature/current influence degradation.

% Cycle axis for illustration
cycles_plot = (1:1000)';

% Baseline synthetic model parameters (will be perturbed per battery later)
a0 = 2e-7;   % Base degradation rate (tuned so capacity stays realistic)
b0 = 2;      % Nonlinear exponent (cycle effect)
c0 = 0.2;    % Temperature influence factor
d0 = 0.3;    % Current influence factor

nominal_temp    = 30; % °C
nominal_current = 3;  % A

% Nominal degradation curve
deg_nom = a0 * cycles_plot.^b0 .* (1 + c0*(nominal_temp-30)/10) .* (1 + d0*(nominal_current-3)/2);
cap_nom = 1 - deg_nom;

figure;
plot(cycles_plot, cap_nom, 'LineWidth', 2);
xlabel('Cycle Number'); ylabel('Capacity (normalized)');
title('Nominal Capacity Fade (30°C, 3A)');
grid on;

% Temperature sweep (keep current fixed)
temps = [20, 30, 40];
figure; hold on;
for i = 1:numel(temps)
    degT = a0 * cycles_plot.^b0 .* (1 + c0*(temps(i)-30)/10) .* (1 + d0*(nominal_current-3)/2);
    capT = 1 - degT;
    plot(cycles_plot, capT, 'LineWidth', 2, 'DisplayName', sprintf('Temp = %d°C', temps(i)));
end
xlabel('Cycle Number'); ylabel('Capacity');
title('Effect of Temperature (Current fixed at 3A)');
legend('Location','best'); grid on;

% Current sweep (keep temperature fixed)
currents = [1, 3, 5];
figure; hold on;
for i = 1:numel(currents)
    degI = a0 * cycles_plot.^b0 .* (1 + c0*(nominal_temp-30)/10) .* (1 + d0*(currents(i)-3)/2);
    capI = 1 - degI;
    plot(cycles_plot, capI, 'LineWidth', 2, 'DisplayName', sprintf('Current = %d A', currents(i)));
end
xlabel('Cycle Number'); ylabel('Capacity');
title('Effect of Current (Temperature fixed at 30°C)');
legend('Location','best'); grid on;

%
%% ==============================================
% PART B) Synthetic Dataset Generation (MULTI-BATTERY) =============================================== We now generate multiple batteries, each with its own operating conditions and slightly different degradation parameters.
% This is crucial for "realistic prognostics": - train on many batteries (full trajectories) - test on a new/unseen battery

% ------------------------------
% Teacher-friendly "knobs"
% ------------------------------
numBatteries = 30;   % number of independent batteries (increase for more data)
N            = 1000; % cycles per battery (length of each trajectory)

% Parameter variability (battery-to-battery differences)
% Example: a0*(1 + 0.20*randn) means ~20% std variability.
sigma_a = 0.20;   % variability on a
sigma_b = 0.03;   % variability on b
sigma_c = 0.10;   % variability on c
sigma_d = 0.10;   % variability on d

% Measurement noise (sensor noise on capacity)
measNoiseStd = 0.003;

% For realism, we keep temperature and current within practical ranges
tempMin = 20; tempMax = 40;   % °C
curMin  = 1;  curMax  = 5;    % A

% Storage for each battery trajectory
batt(numBatteries) = struct();

for k = 1:numBatteries

    cycles = (1:N)';

    % ---- Battery-specific parameter perturbation ----
    a = max(a0 * (1 + sigma_a*randn), 1e-9);  % ensure positive
    b = b0 + sigma_b*randn;
    c = c0 * (1 + sigma_c*randn);
    d = d0 * (1 + sigma_d*randn);

    % ---- Operating conditions as slowly varying signals (more realistic than i.i.d. noise) ----
    % We create a sinusoidal baseline + random noise, then clamp to bounds.
    tempPhase = 2*pi*rand;
    curPhase  = 2*pi*rand;

    temperature = 30 + 5*sin(2*pi*cycles/250 + tempPhase) + 1.5*randn(N,1);
    current     = 3  + 1*sin(2*pi*cycles/180 + curPhase)  + 0.4*randn(N,1);

    temperature = min(max(temperature, tempMin), tempMax);
    current     = min(max(current, curMin),  curMax);

    % ---- Nonstationary multiplicative "process noise" (aging uncertainty) ----
    % Larger uncertainty at later life + small oscillatory effect
    noiseFactor = 1 + 0.03*sin(0.02*cycles + 2*pi*rand) + 0.05*(cycles/N).*randn(N,1);
    noiseFactor = max(noiseFactor, 0.7); % keep positive / avoid extreme negatives

    % ---- Degradation model (synthetic) ----
    degradation = a * cycles.^b ...
        .* (1 + c*(temperature-30)/10) ...
        .* (1 + d*(current-3)/2) ...
        .* noiseFactor;

    capacity = 1 - degradation;

    % Add small measurement noise (sensor noise)
    capacity = capacity + measNoiseStd*randn(N,1);

    % Optional: clamp to physical range [0,1]
    capacity = min(max(capacity, 0), 1);

    % Cycle normalized feature (0..1) helps when batteries have different lengths
    cycleNorm = cycles / N;

    % Store battery data
    batt(k).id          = k;
    batt(k).cycles      = cycles;
    batt(k).cycleNorm   = cycleNorm;
    batt(k).temperature = temperature;
    batt(k).current     = current;
    batt(k).capacity    = capacity;

    % Feature matrix for time-series modeling
    % Each row = one time step (cycle), each column = one feature
    % We include capacity itself as an input feature because in prognostics
    % we typically measure current capacity and want to forecast future capacity.
    batt(k).X = [cycleNorm, temperature, current, capacity]; % N x 4
    batt(k).Y = capacity;                                    % N x 1
end

% Visualize a few battery trajectories (capacity vs cycles)
figure; hold on;
plotCount = min(10, numBatteries);
for k = 1:plotCount
    plot(batt(k).cycles, batt(k).capacity, 'LineWidth', 1.5, ...
        'DisplayName', sprintf('Battery %d', k));
end
xlabel('Cycle Number'); ylabel('Capacity');
title('Example Capacity Trajectories (Multiple Batteries)');
legend('Location','best'); grid on;

%
%% =====================================================
% PART C) Supervised Dataset via Sliding Windows (by battery) ====================================================== We now convert each battery time series into MANY training samples.
% Sliding window idea: - Choose a window length W (lookback). - Choose a horizon L (L-step-ahead forecast).
% For each window starting at index s: Input sequence X(s : s+W-1) Target sequence Y(s+L : s+L+W-1)
% Because the LSTM outputs a sequence (OutputMode='sequence'),
% the network learns: at each time step in the input, predict L steps ahead. The last output element corresponds to "predict L steps ahead from the end of the window".

% ------------------------------
% Teacher-friendly "knobs"
% ------------------------------
windowLength = 100;   % W = lookback length (how much history we show the network)
horizon      = 20;    % L = how many steps ahead to predict (forecast horizon)

% ------------------------------
% Split by battery ID (REALISTIC!)
% ------------------------------
% Train/Validation/Test split: batteries are disjoint
trainFrac = 0.70;
valFrac   = 0.15;
testFrac  = 0.15;

perm = randperm(numBatteries);
nTrainB = round(trainFrac*numBatteries);
nValB   = round(valFrac*numBatteries);
nTestB  = numBatteries - nTrainB - nValB;

trainIDs = perm(1:nTrainB);
valIDs   = perm(nTrainB+1 : nTrainB+nValB);
testIDs  = perm(nTrainB+nValB+1 : end);

fprintf('\nBattery split (by ID):\n');
fprintf('  Train batteries: %d\n', numel(trainIDs));
fprintf('  Val batteries  : %d\n', numel(valIDs));
fprintf('  Test batteries : %d\n\n', numel(testIDs));

% Build sliding-window datasets
[XTrain, YTrain, infoTrain] = buildSlidingWindowDataset(batt, trainIDs, windowLength, horizon);
[XVal,   YVal,   infoVal  ] = buildSlidingWindowDataset(batt, valIDs,   windowLength, horizon);
[XTest,  YTest,  infoTest ] = buildSlidingWindowDataset(batt, testIDs,  windowLength, horizon);

fprintf('Number of sliding-window sequences:\n');
fprintf('  Train: %d\n', numel(XTrain));
fprintf('  Val  : %d\n', numel(XVal));
fprintf('  Test : %d\n\n', numel(XTest));

% Show one example to students: format expected by trainNetwork for sequences
% XTrain{1} is [numFeatures x windowLength]
fprintf('Example sequence sizes:\n');
fprintf('  size(XTrain{1}) = %d x %d  (features x time)\n', size(XTrain{1},1), size(XTrain{1},2));
fprintf('  size(YTrain{1}) = %d x %d  (responses x time)\n\n', size(YTrain{1},1), size(YTrain{1},2));

%
%% =====================================================
% Normalization (IMPORTANT for neural networks) ====================================================== We compute normalization statistics ONLY from TRAINING batteries to avoid leaking test information.

% Collect raw feature rows and target values from training batteries
Xall = [];
Yall = [];
for id = trainIDs
    Xall = [Xall; batt(id).X]; %#ok<AGROW>  N x numFeatures
    Yall = [Yall; batt(id).Y]; %#ok<AGROW>  N x 1
end

muX    = mean(Xall, 1)';      % numFeatures x 1
sigmaX = std(Xall, 0, 1)';    % numFeatures x 1
sigmaX(sigmaX < 1e-12) = 1;   % avoid divide-by-zero

muY    = mean(Yall);          % scalar
sigmaY = std(Yall);
if sigmaY < 1e-12
    sigmaY = 1;
end

% Normalize the cell array sequences
XTrainN = normalizeCellSeq(XTrain, muX, sigmaX);
XValN   = normalizeCellSeq(XVal,   muX, sigmaX);
XTestN  = normalizeCellSeq(XTest,  muX, sigmaX);

YTrainN = normalizeCellSeq(YTrain, muY, sigmaY);
YValN   = normalizeCellSeq(YVal,   muY, sigmaY);
YTestN  = normalizeCellSeq(YTest,  muY, sigmaY);

%
%% =====================================================
% Define Deep LSTM Network ======================================================

numFeatures  = size(XTrainN{1},1);
numResponses = size(YTrainN{1},1); % for us: 1 (capacity)

numHiddenUnits = 10;

layers = [ ...
    sequenceInputLayer(numFeatures, 'Name','input')

    lstmLayer(numHiddenUnits, 'OutputMode','sequence', 'Name','lstm1')
    fullyConnectedLayer(50, 'Name','fc1')
    dropoutLayer(0.5, 'Name','drop1')

    fullyConnectedLayer(numResponses, 'Name','fc_out')
    regressionLayer('Name','regression') ...
];

% Visualize network (nice in class)
% Visualize network (nice in class)
if exist('analyzeNetwork','file')
    analyzeNetwork(layers);
else
    fprintf('analyzeNetwork not available in this MATLAB version. Skipping network visualization.\n');
end

%
%% =====================================================
% Training Options (tunable hyperparameters) ======================================================

maxEpochs   = 5;
miniBatchSz = 64;
learnRate   = 1e-2;   % LSTMs usually behave better with 1e-3 than 1e-2

% ------------------------------
% Training options
% ------------------------------




options = trainingOptions('sgdm', ...
    'MaxEpochs', maxEpochs, ...
    'MiniBatchSize', miniBatchSz, ...
    'InitialLearnRate', learnRate, ...
    'GradientThreshold', 1, ...
    'Shuffle','every-epoch', ...
    'ValidationData',{XValN, YValN}, ...
    'ValidationFrequency', 30, ...
    'Verbose', true, ...
    'Plots','training-progress');

%
%% =====================================================
% Train the network ======================================================

fprintf('\nTraining deep LSTM...\n');
net = trainNetwork(XTrainN, YTrainN, layers, options);

%
%% =====================================================
% Evaluate on Test Batteries (L-step-ahead forecast points) ======================================================
% Predict normalized outputs

YPredN = predict(net, XTestN, 'MiniBatchSize', miniBatchSz);

% Denormalize predictions back to capacity scale
YPred = denormalizeCellSeq(YPredN, muY, sigmaY);

% Extract the "forecast point" = last element in each target sequence
% This corresponds to predicting horizon steps ahead from the END of the window.
yTruePoint = extractLastElement(YTest);
yPredPoint = extractLastElement(YPred);

% Compute global metrics on test set
rmseTest = sqrt(mean((yPredPoint - yTruePoint).^2));
maeTest  = mean(abs(yPredPoint - yTruePoint));

fprintf('\n--- Test Performance (forecast point = last element of each Y window) ---\n');
fprintf('RMSE = %.6f\n', rmseTest);
fprintf('MAE  = %.6f\n', maeTest);

% ------------------------------
% Baseline for teaching: "Persistence"
% ------------------------------
% Baseline idea: predict future capacity as "same as last observed capacity"
% in the input window (i.e., assume no change).
idxCapFeature = 4; % our features: [cycleNorm, temperature, current, capacity]
yBase = zeros(numel(XTest),1);
for i = 1:numel(XTest)
    yBase(i) = XTest{i}(idxCapFeature, end); % last observed capacity in input
end
rmseBase = sqrt(mean((yBase - yTruePoint).^2));
maeBase  = mean(abs(yBase - yTruePoint));

fprintf('\nBaseline (persistence: y(t+L) ≈ y(t))\n');
fprintf('RMSE = %.6f\n', rmseBase);
fprintf('MAE  = %.6f\n', maeBase);

%
%% =====================================================
% Plot results for ONE test battery (trajectory style) ======================================================

testBatteryToPlot = testIDs(1);

mask = (infoTest.batteryId == testBatteryToPlot);
cyc  = infoTest.forecastCycle(mask);
yt   = yTruePoint(mask);
yp   = yPredPoint(mask);

% Sort by cycle (should already be sorted, but safe)
[cyc, order] = sort(cyc);
yt = yt(order);
yp = yp(order);

figure;
plot(cyc, yt, 'LineWidth', 2, 'DisplayName','True');
hold on;
plot(cyc, yp, '--', 'LineWidth', 2, 'DisplayName', sprintf('Predicted (L=%d)', horizon));
xlabel('Cycle Number (forecast point)');
ylabel('Capacity');
title(sprintf('LSTM Forecast on Unseen Test Battery #%d (L=%d steps ahead)', testBatteryToPlot, horizon));
legend('Location','best'); grid on;

% Error plot
figure;
plot(cyc, yp-yt, 'LineWidth', 1.5);
xlabel('Cycle Number (forecast point)');
ylabel('Prediction Error (Pred - True)');
title(sprintf('Forecast Error on Test Battery #%d', testBatteryToPlot));
grid on;

% Scatter plot for teaching
figure;
scatter(yt, yp, 20, 'filled');
xlabel('True Capacity'); ylabel('Predicted Capacity');
title('Test Forecast Points: True vs Predicted');
grid on; axis equal;

%
%% =====================================================
% (Optional) Simple End-of-Life (EOL) / RUL demonstration ====================================================== A common prognostics task is to predict when capacity crosses a threshold. For example, EOL threshold = 0.70 (70% of nominal).

eolThresh = 0.70;

trueEOL = find(yt <= eolThresh, 1, 'first');
predEOL = find(yp <= eolThresh, 1, 'first');

if ~isempty(trueEOL)
    trueEOLcycle = cyc(trueEOL);
else
    trueEOLcycle = NaN;
end

if ~isempty(predEOL)
    predEOLcycle = cyc(predEOL);
else
    predEOLcycle = NaN;
end

fprintf('\n--- Simple EOL demo (Battery #%d, threshold = %.2f) ---\n', testBatteryToPlot, eolThresh);
fprintf('True EOL cycle (first <= thresh): %g\n', trueEOLcycle);
fprintf('Pred EOL cycle (first <= thresh): %g\n', predEOLcycle);

figure;
plot(cyc, yt, 'LineWidth', 2, 'DisplayName','True');
hold on;
plot(cyc, yp, '--', 'LineWidth', 2, 'DisplayName','Predicted');
yline(eolThresh, ':', 'LineWidth', 2, 'DisplayName', sprintf('EOL threshold %.2f', eolThresh));
if ~isnan(trueEOLcycle); xline(trueEOLcycle, '-',  'LineWidth', 1.5, 'DisplayName','True EOL'); end
if ~isnan(predEOLcycle); xline(predEOLcycle, '--', 'LineWidth', 1.5, 'DisplayName','Pred EOL'); end
xlabel('Cycle Number (forecast point)');
ylabel('Capacity');
title(sprintf('EOL Threshold Crossing (Battery #%d)', testBatteryToPlot));
legend('Location','best'); grid on;

%

%% ================================================================
%  QUESTIONS (Concept check)
%  Goal: Ensure you understand what the script is doing and *why*.
%  Answer in 2–5 lines each, with an equation or a short diagram when relevant.
%  ================================================================

% Q1) Why do we generate MANY batteries instead of training on one battery only?
%     Explain in terms of generalization and "training range" / extrapolation failure.

% 

% Q3) What is the shape of one training sample (for sequence input)?
%     Give the dimensions of XTrain{i} and YTrain{i}.
%     Why are sequences stored in CELL arrays for trainNetwork?

% Q4) We split the dataset by BATTERY ID (not by time).
%     What would go wrong if we split by time on a single battery, and why?

% 

% Q7) Interpret the training curves:
%     If training RMSE decreases but validation RMSE increases, what does it indicate?
%     Give two changes you would try in the network/training settings.



%


%% ================================================================
%  EXERCISES (Hands-on trials)
%  Each exercise asks you to modify 1–2 lines only.
%  Record your results in a small table:
%     (W, L, hiddenUnits, dropout, learnRate, TestRMSE, TestMAE, Notes)
%  ================================================================

% EX1) Sliding window sensitivity (W and L)
% Trial steps:
%   (a) Fix horizon L = 5. Try windowLength W = 20, 50, 100.
%   (b) Fix windowLength W = 50. Try horizon L = 1, 5, 10, 20.
% What to observe:
%   - Larger W can help capture slower trends, but may increase training time.
%   - Larger L makes forecasting harder; RMSE should generally increase with L.
% Evaluation:
%   - You should produce a plot RMSE vs L and briefly interpret it.

% EX2) Network depth and capacity
% Trial steps:
%   (a) Use 1 LSTM block (remove two LSTM blocks).
%   (b) Use 2 blocks, then 3 blocks (original).
%   (c) Try numHiddenUnits = 50 vs 100 vs 200.
% What to observe:
%   - Too small: underfitting (both train and val errors high).
%   - Too large: overfitting (train low, val high) unless regularized.
% Evaluation:
%   - Report which configuration gives best validation RMSE and why.

% EX3) Regularization via dropout
% Trial steps:
%   Keep architecture fixed, change dropout p = 0.0, 0.1, 0.2, 0.5.
% What to observe:
%   - p too high can cause underfitting (model learns "flat" predictions).
% Evaluation:
%   - Explain the bias–variance trade-off you observe.

% EX4) Learning rate and convergence
% Trial steps:
%   Run with learnRate = 1e-2, 1e-3, 3e-4.
% What to observe:
%   - Too high: unstable training/validation curves.
%   - Too low: slow learning, may not converge within maxEpochs.
% Evaluation:
%   - Identify a learning rate that gives stable training and good validation RMSE.

% 

% EX6) Train/Val/Test split by battery (data generalization)
% Trial steps:
%   Change the number of training batteries:
%       Train = 60%, 70%, 80% of batteries (keep test batteries fixed).
% What to observe:
%   - More training batteries usually improves generalization to unseen batteries.
% Evaluation:
%   - Provide a short explanation relating dataset diversity to generalization.





%% =========================================================================
% LOCAL HELPER FUNCTIONS (kept at end for readability) =========================================================================

function [XCell, YCell, info] = buildSlidingWindowDataset(batt, batteryIDs, W, L)
% buildSlidingWindowDataset
% -------------------------
% Convert multiple battery trajectories into a set of supervised sequence pairs:
%   XCell{i} = [numFeatures x W] input sequence
%   YCell{i} = [numResponses x W] target sequence (capacity shifted by L)
%
% Inputs
%   batt      : struct array with fields batt(k).X (N x numFeatures), batt(k).Y (N x 1)
%   batteryIDs: vector of battery indices to include
%   W         : window length (lookback)
%   L         : horizon (L-step ahead)
%
% Outputs
%   XCell, YCell : cell arrays of sequences
%   info         : struct with per-sample metadata (batteryId, forecastCycle)

    XCell = {};
    YCell = {};

    batteryIdVec   = [];
    forecastCycleV = [];

    for id = batteryIDs

        X = batt(id).X;          % N x numFeatures
        Y = batt(id).Y;          % N x 1
        cycles = batt(id).cycles;

        N = size(X,1);

        % Number of windows possible:
        % start index s must satisfy: s + W - 1 + L <= N
        lastStart = N - W - L + 1;

        for s = 1:lastStart

            idxX = s : (s+W-1);
            idxY = (s+L) : (s+L+W-1);

            % trainNetwork expects: features x time
            Xseq = X(idxX, :)';   % [numFeatures x W]
            Yseq = Y(idxY, :)';   % [1 x W]

            XCell{end+1,1} = Xseq; %#ok<AGROW>
            YCell{end+1,1} = Yseq; %#ok<AGROW>

            % The forecast point we often evaluate is the last target element,
            % which corresponds to cycle index: (s+W-1) + L
            forecastCycle = cycles(s+W-1+L);

            batteryIdVec(end+1,1)   = id;           %#ok<AGROW>
            forecastCycleV(end+1,1) = forecastCycle; %#ok<AGROW>
        end
    end

    info = struct();
    info.batteryId    = batteryIdVec;
    info.forecastCycle = forecastCycleV;
end

function XN = normalizeCellSeq(XCell, mu, sigma)
% normalizeCellSeq
% ----------------
% Apply z-score normalization to each sequence in a cell array:
%   XN = (X - mu) ./ sigma
%
% Works for both:
%  - mu/sigma vectors (features) for X
%  - mu/sigma scalars for Y

    XN = XCell;
    for i = 1:numel(XCell)
        X = XCell{i};
        XN{i} = (X - mu) ./ sigma;
    end
end

function X = denormalizeCellSeq(XNCell, mu, sigma)
% denormalizeCellSeq
% ------------------
% Undo z-score normalization:
%   X = XN .* sigma + mu

    X = XNCell;
    for i = 1:numel(XNCell)
        XN = XNCell{i};
        X{i} = XN .* sigma + mu;
    end
end

function yLast = extractLastElement(YCell)
% extractLastElement
% ------------------
% For each target sequence YCell{i} (1 x W), return the last element.
% This corresponds to "forecast point" for horizon L from the end of the window.

    yLast = zeros(numel(YCell),1);
    for i = 1:numel(YCell)
        yi = YCell{i};
        yLast(i) = yi(end);
    end
end