% DTS_IsothermSlope_FieldTest1.m
%
% Isotherm-slope flow profiling applied to measured DTS from Utah FORGE well
% 16B(78)-32, huff-puff Test 1. Companion code for Becker and Ifrene,
% Geothermics.
%
% Produces two figures:
%   Figure 2  DTS waterfall for Test 1 with the 205 C isotherm emphasized,
%             stage boundaries, and the shaded disturbed-zone (slope-fit) mask.
%   Figure 5  Three panels: injection-phase isotherm tracking, production-phase
%             isotherm tracking, and the derived flow profile Q(z)/Q_ref.
% and the injection and production flow balances reported in Table 2, plus the
% reference-interval retardation ratios quoted in Section 3.3.
%
% Method. The tracked isotherm is extracted from the full 2-D DTS temperature
% field with MATLAB contourc() and a phase-specific envelope selection:
%   injection  = deepest T_iso crossing at each time step
%   production = shallowest T_iso crossing at each time step
% Straight lines are then fitted to the tracked front within manually defined
% depth windows above and below each stage, and the ratio of fitted slopes is
% converted into flow partitioning. Production stages may be gains or losses;
% fitted velocity ratios are not clamped to enforce one-way flow.
%
% MASK_Z_INJ / MASK_Z_PROD are slope-fit masks only. They do not alter contour
% tracking. Selected-front points falling inside those intervals are excluded
% from polyfit and are plotted in gray (Figure 5).
%
% Depths are held internally in feet (the units of the DTS acquisition file)
% and relabelled to metres for display via set_depth_axis.
%
% Requires FORGE_16B_Test1_DTS.mat in the same folder as this script. That
% file holds only the Test 1 window of the calibrated DTS record, already in
% deg C, local site time, and ft MD. It is produced from the full campaign
% record by make_Test1_DTS_mat.m, which is not part of this supplement.

clc; clearvars;

% =========================================================
% ---- EDIT THESE ONLY ----
% =========================================================
% Every parameter below is documented with its units, how it was chosen for
% Test 1, and what goes wrong if it is set badly. Nothing below the
% "Wellbore geometry" heading needs to be changed to reproduce the article
% figures.

% ---- Input file -----------------------------------------------------------
% FILE_DTS holds only the Test 1 window, already in deg C, local site time,
% and ft MD. See make_Test1_DTS_mat.m for how it is produced.
DATA_FOLDER    = fileparts(mfilename('fullpath'));
FILE_DTS       = 'FORGE_16B_Test1_DTS.mat';
DTS_TEMP_UNITS = 'C';    % colorbar label only; no conversion is applied

% ---- Record gaps ----------------------------------------------------------
% Acquisition breaks are blanked to NaN before tracking, so the isotherm
% tracker cannot draw a straight line across a period with no data and report
% it as a slope. A break is flagged when the interval between two consecutive
% samples exceeds
%       threshold = max(DTS_GAP_FACTOR * median sample interval,
%                       DTS_MIN_GAP_MINUTES)
% The script prints the effective threshold and which of the two terms binds.
%
% DTS_GAP_FACTOR (dimensionless) keeps ordinary jitter in acquisition timing
%   from being flagged. Lower it if real dropouts are being missed; raise it
%   if normal timing scatter is being blanked.
% DTS_MIN_GAP_MINUTES (minutes) is an absolute floor. At DTS sampling rates
%   of seconds, DTS_GAP_FACTOR alone would flag pauses of only a few minutes,
%   so this floor is usually the term that actually governs.
DTS_GAP_FACTOR      = 8.0;
DTS_MIN_GAP_MINUTES = 20.0;

% ---- Isotherm tracking ----------------------------------------------------
% CONTOUR_MAX_FILL_GAP_STEPS (number of TIME STEPS, not minutes) bridges
% isolated time steps where the contour tracker found no crossing, by linear
% interpolation. 2 bridges only one or two dropped samples. Set 0 to disable.
% Keep it small: a large value re-introduces exactly the interpolation across
% missing data that the gap blanking above exists to prevent.
CONTOUR_MAX_FILL_GAP_STEPS = 2;

% Z_TRACK_MIN / Z_TRACK_MAX (ft MD) bound the depth band the contour tracker
% searches, and disambiguate which crossing is selected where the isotherm is
% multivalued. This is NOT a plot limit. It must bracket every fit window
% defined below. Too narrow truncates the tracked front; too wide admits
% crossings from outside the monitored interval, such as cable-exit artifacts
% below the completion.
Z_TRACK_MIN = 8500;
Z_TRACK_MAX = 10000;

% ---- Well completion ------------------------------------------------------
% Perforation intervals by stage, ft MD, one [top bottom] row per cluster.
% Used for the stage boundary lines on the figures and for the stage ordering
% in the flow balance. Changing these does not move any fit window.
FRAC_STAGES.Stage1 = [9769,9773; 9756,9760; 9745,9749; 9690,9694];
FRAC_STAGES.Stage2 = [9508,9512; 9475,9479; 9459,9463; 9447,9451; 9429,9433];
FRAC_STAGES.Stage3 = [9389,9393; 9343,9347; 9265,9269];
FRAC_STAGES.Stage4 = [9054,9058; 9026,9030; 8995,8999; 8958,8962];
FRAC_STAGES.Stage5 = [8879,8883; 8870,8874; 8834,8838; 8774,8778];

% REF_Z_MIN / REF_Z_MAX (ft MD) delimit the interval used for the apparent
% retardation ratio R = v_piston / v_isotherm reported in Section 3.3. The
% interval must contain NO fluid entry or exit, so that the isotherm there
% moves at the retarded thermal-front speed alone. Here it runs from the top
% of the tracked band to just above the shallowest Stage 5 perforation
% (8774 ft). Any exchange inside this interval biases R.
%
% DTS_Synthetic_PDE_Test1_GainLoss.m computes its own retardation over the
% same interval, and RHO_CP_TUNE there is calibrated against the value printed
% here. If these depths are changed, change them in that script too, or the
% two ratios are no longer computed over the same depths.
REF_Z_MIN = 8500;
REF_Z_MAX = 8760;

% ---- Test 1 events and metered rates --------------------------------------
% Event times in local site time, matching make_Test1_DTS_mat.m. All plotted
% times are minutes elapsed from INJECTION_TIME.
INJECTION_TIME  = datetime(2025,8,26,13,38,0);
PRODUCTION_TIME = datetime(2025,8,26,17, 5,0);
SHUTIN_TIME     = datetime(2025,8,26,21,20,0);

% Metered surface rates, barrels per minute. These set the absolute scale of
% the flow profile and the m^3/hr labels in Figure 5. The partitioning
% PERCENTAGES do not depend on them: the method uses ratios of fitted slopes,
% anchored to the surface rate only at the top of the profile.
Q_inj_bpm  = 2.5;      % 23.8 m^3/hr
Q_prod_bpm = 2.47;     % 23.6 m^3/hr

% T_iso (deg C) is the tracked isotherm. It must remain continuously
% trackable across every fit window below; an isotherm that runs out of the
% depth range partway through the test cannot be fitted above and below a
% stage. Absolute isotherm velocity depends on which isotherm is chosen
% (cooler isotherms propagate more slowly), but the flow partitioning is
% first-order insensitive to that choice because it uses velocity ratios.
T_iso = 205;

% ---- Disturbed-zone masks -------------------------------------------------
% [n x 2] arrays of [z_top z_bottom] in ft MD, or [] to disable. Points whose
% depth falls inside these intervals are excluded from the slope fits and
% drawn in gray on Figure 5. They do NOT alter contour tracking; the front is
% still tracked through them.
% Picked by eye from the waterfall (Figure 2) where the isotherm is near
% vertical or jagged with excursions of both signs, indicating local
% circulation across a stage rather than a net change in throughflow
% (Section 2.4). Fitting such an interval would return a near-piston front
% velocity and a spurious flow change.
MASK_Z_INJ = [
                9280 9360
                9460 9500
            ];
MASK_Z_PROD = [9320 9490];

% ---- Slope-fit windows ----------------------------------------------------
% Depth windows, ft MD, in the convention
%       [shallow_min shallow_max deep_min deep_max]
% for both phases. Each pair is sorted internally, so order within a pair does
% not matter. The shallow pair brackets the front above the stage's exchange
% point and the deep pair brackets it below.
%
% Choosing a window:
%   - it must lie wholly on one side of that stage's slope break, never
%     straddle it, or the fit averages two velocities and washes out the step;
%   - it must not reach into a neighbouring stage's break;
%   - it needs at least 3 tracked points, otherwise the stage is skipped with
%     a message and simply passes flow through unchanged;
%   - the fitted breakpoint depth is the INTERSECTION of the two lines
%     (z_cross), not the stage midpoint, so a window does not have to be
%     centred on the perforations.
%
% INJ_ON / PROD_ON (logical) select which stages enter the flow balance.
% Setting one false leaves that stage in the figure but carries flow through
% it unchanged, with zero net exchange.
INJ_FIT.Stage5 = [8750, 8850, 8830, 8930]; INJ_ON.Stage5 = true;
INJ_FIT.Stage4 = [8950, 9050, 9100, 9200]; INJ_ON.Stage4 = true;
INJ_FIT.Stage3 = [9100, 9200, 9370, 9470]; INJ_ON.Stage3 = true;
INJ_FIT.Stage2 = [9370, 9470, 9490, 9580]; INJ_ON.Stage2 = true;
INJ_FIT.Stage1 = [9600, 9680, 9720, 9800]; INJ_ON.Stage1 = true;

% Stage 3 production is disabled: no resolvable production-phase slope break
% (Section 3.2). Its window is retained so the choice can be re-examined.
PROD_FIT.Stage5 = [8550, 8756, 8800, 8930]; PROD_ON.Stage5 = true;
PROD_FIT.Stage4 = [8950, 9050, 9100, 9300]; PROD_ON.Stage4 = true;
PROD_FIT.Stage3 = [9110, 9250, 9300, 9420]; PROD_ON.Stage3 = false;
PROD_FIT.Stage2 = [9220, 9535, 9560, 9670]; PROD_ON.Stage2 = true;
PROD_FIT.Stage1 = [9560, 9670, 9670, 9850]; PROD_ON.Stage1 = true;

% =========================================================
% ---- DERIVED: no editing needed below this line ----
% =========================================================

% Phase time windows, minutes from injection start. The moving_down and
% moving_up masks built below already restrict fitting to each phase.
INJ_T_MIN  = 0;
INJ_T_MAX  = minutes(PRODUCTION_TIME - INJECTION_TIME);   % injection ends
PROD_T_MIN = INJ_T_MAX;                                   % production begins
PROD_T_MAX = minutes(SHUTIN_TIME     - INJECTION_TIME);   % shut-in

% Active stage lists from the enable flags above.
inj_order       = {'Stage5','Stage4','Stage3','Stage2','Stage1'};   % shallow -> deep
fit_stages_inj  = inj_order(cellfun(@(s) INJ_ON.(s),  inj_order));
prod_order      = {'Stage1','Stage2','Stage3','Stage4','Stage5'};   % deep -> shallow
fit_stages_prod = prod_order(cellfun(@(s) PROD_ON.(s), prod_order));

% =========================
% Wellbore geometry
% =========================
% Flow occurs in the casing-CT annulus (coiled tubing is run inside casing for
% the full stage interval), so A_flow, not the full bore, sets the piston
% velocity. This cross-section does NOT enter the flow partitioning below,
% which uses velocity ratios and the metered rate only; A_flow is used solely
% for the piston velocity and the apparent-retardation diagnostic.
r_casing = (5.920/2) * 0.0254;     % m, casing ID
A_casing = pi * r_casing^2;        % m^2, full bore
OD_CT    = 1.75 * 0.0254;          % m, coiled tubing OD (Leggett et al., 2026)
A_CT     = pi * (OD_CT/2)^2;       % m^2
A_flow   = A_casing - A_CT;        % m^2, annulus flow area

% =========================
% Load the Test 1 DTS record
% =========================
% The .mat is already conditioned: temperature in deg C, time in local site
% time, depth in ft MD and increasing. No unit or clock conversion here.
path_dts = fullfile(DATA_FOLDER, FILE_DTS);
if ~isfile(path_dts)
    error(['DTS file not found:\n  %s\n' ...
           'Place %s in the same folder as this script, or edit ' ...
           'DATA_FOLDER / FILE_DTS above.'], path_dts, FILE_DTS);
end
S = load(path_dts);
required = {'z_ft','t_local','T_C'};
missing  = required(~isfield(S, required));
if ~isempty(missing)
    error('%s is missing variable(s): %s', FILE_DTS, strjoin(missing, ', '));
end
z_dts   = double(S.z_ft(:));
t_dts_w = S.t_local(:);
T_w     = double(S.T_C);
T_w(~isfinite(T_w)) = NaN;
if isfield(S,'meta')
    fprintf('Loaded %s: %s, %s to %s\n', FILE_DTS, S.meta.test, ...
        datestr(S.meta.window_start), datestr(S.meta.window_end));
end
clear S

% Blank any break in the record so the tracker cannot interpolate across it.
% The effective threshold is reported, along with which of the two settings
% actually binds, since at typical DTS sampling rates it is usually the floor.
[gap_intervals, dt_nom, gap_thr] = detect_time_gaps( ...
    t_dts_w, DTS_GAP_FACTOR, DTS_MIN_GAP_MINUTES);
if DTS_GAP_FACTOR*dt_nom >= DTS_MIN_GAP_MINUTES
    gap_binds = 'DTS_GAP_FACTOR';
else
    gap_binds = 'DTS_MIN_GAP_MINUTES';
end
fprintf(['Sample interval %.2f min; gap threshold %.1f min (set by %s); ' ...
         '%d gap(s) blanked\n'], dt_nom, gap_thr, gap_binds, size(gap_intervals,1));
if ~isempty(gap_intervals)
    T_w(mask_rows_from_gap_intervals(t_dts_w, gap_intervals), :) = NaN;
end

% Crop to the Test 1 window, injection start to shut-in.
mask_test = t_dts_w >= INJECTION_TIME & t_dts_w <= SHUTIN_TIME;
t_test    = t_dts_w(mask_test);
T_test    = T_w(mask_test, :);           % [ntime x ndepth]
t_min     = minutes(t_test - INJECTION_TIME);

fprintf('Test 1 window: %s to %s (%d time steps, %d depths)\n', ...
    datestr(INJECTION_TIME), datestr(SHUTIN_TIME), numel(t_test), numel(z_dts));

% Piston velocities from the metered rate and the annulus area.
v_piston_inj  = (Q_inj_bpm  * 0.158987 / A_flow) * 3.28084;   % ft/min
v_piston_prod = (Q_prod_bpm * 0.158987 / A_flow) * 3.28084;   % ft/min

% Unit conversions for figure display only. Internal calculations stay in bpm
% and ft/min and are ratio-based, so units cancel.
BPM_TO_M3HR   = 0.158987 * 60;   % bpm    -> m^3/hr  (= 9.53922)
FTMIN_TO_MMIN = 0.3048;          % ft/min -> m/min

%% =========================================================
%           INJECTION isotherm slope analysis
% =========================================================
z_iso_inj = contour_front_intersections( ...
    t_min, z_dts, T_test, T_iso, [Z_TRACK_MIN Z_TRACK_MAX], [INJ_T_MIN INJ_T_MAX], +1, ...
    CONTOUR_MAX_FILL_GAP_STEPS);
z_iso_inj   = z_iso_inj(:);
moving_down = isfinite(z_iso_inj) & t_min(:)>=INJ_T_MIN & t_min(:)<INJ_T_MAX;

% Points feeding the slope fits: the tracked front within the injection phase,
% less any point inside a disturbed-zone mask.
[tf_inj, zf_inj] = build_fit_points(t_min, z_iso_inj, moving_down);
mask_slopefit_inj = points_in_depth_mask(zf_inj, MASK_Z_INJ);
tf_inj = tf_inj(~mask_slopefit_inj);  zf_inj = zf_inj(~mask_slopefit_inj);
% Same mask indexed against the full tracked front, for gray plotting.
mask_slope_inj = points_in_depth_mask(z_iso_inj, MASK_Z_INJ);

nS = numel(fit_stages_inj);
v_shallow_inj = nan(nS,1);  v_deep_inj = nan(nS,1);
Ql_Qu_inj = nan(nS,1);      Qf_Qu_inj  = nan(nS,1);
z_cross_inj = nan(nS,1);    t_cross_inj = nan(nS,1);
p_shallow_inj_all = nan(nS,2);  p_deep_inj_all = nan(nS,2);   % stored for plotting

fprintf('\n=== INJECTION isotherm slope analysis (T_iso=%.0f C, Qinj=%.2f bpm) ===\n', T_iso, Q_inj_bpm);
fprintf('Piston velocity (annulus): %.2f ft/min\n', v_piston_inj);
fprintf('%-10s %10s %10s %10s %10s\n', 'Stage','v_shallow','v_deep','Ql/Qu','Qf/Qu');
fprintf('%s\n', repmat('-',1,55));
for si = 1:nS
    lims = INJ_FIT.(fit_stages_inj{si});
    [shallow_min, shallow_max, deep_min, deep_max] = unpack_inj_fit(lims);
    mask_shallow = zf_inj>=shallow_min & zf_inj<=shallow_max;
    mask_deep    = zf_inj>=deep_min    & zf_inj<=deep_max;
    if sum(mask_shallow)<3 || sum(mask_deep)<3
        fprintf('%-10s  insufficient points for fit\n', fit_stages_inj{si}); continue
    end
    p_shallow = polyfit(tf_inj(mask_shallow), zf_inj(mask_shallow), 1);
    p_deep    = polyfit(tf_inj(mask_deep),    zf_inj(mask_deep),    1);
    p_shallow_inj_all(si,:) = p_shallow;  p_deep_inj_all(si,:) = p_deep;
    v_shallow_inj(si) = p_shallow(1);
    v_deep_inj(si)    = p_deep(1);
    Ql_Qu_inj(si)     = v_deep_inj(si) / v_shallow_inj(si);   % fraction continuing deeper
    Qf_Qu_inj(si)     = 1 - Ql_Qu_inj(si);                    % fraction exchanged at this stage
    t_cross_inj(si)   = (p_deep(2) - p_shallow(2)) / (p_shallow(1) - p_deep(1));
    z_cross_inj(si)   = p_shallow(1)*t_cross_inj(si) + p_shallow(2);   % breakpoint depth
    fprintf('%-10s %10.2f %10.2f %10.3f %10.3f\n', ...
        fit_stages_inj{si}, v_shallow_inj(si), v_deep_inj(si), Ql_Qu_inj(si), Qf_Qu_inj(si));
end

% Injection flow balance, anchored at the metered surface rate and walked
% shallow -> deep. Q is the flow remaining in the wellbore below each stage;
% Net < 0 is fluid lost to the formation at that stage.
Qf_bpm_inj = nan(nS,1);  Ql_bpm_inj = nan(nS,1);
Qu_i = Q_inj_bpm;
for si = 1:nS
    if isnan(Qf_Qu_inj(si)), continue; end
    Qf_bpm_inj(si) = Qu_i * Qf_Qu_inj(si);
    Ql_bpm_inj(si) = Qu_i * Ql_Qu_inj(si);
    Qu_i = Ql_bpm_inj(si);
end
stage_order_sd   = flipud(get_stage_order_deep_to_shallow(FRAC_STAGES, fieldnames(FRAC_STAGES)));
Q_below_inj_full = fill_passthrough(stage_order_sd, fit_stages_inj, Ql_bpm_inj, Q_inj_bpm);
T_inj = build_flow_balance(stage_order_sd, Q_below_inj_full, Q_inj_bpm, 'inject', 'Barefoot');
fprintf('\n--- INJECTION flow balance (anchored at surface Qinj=%.3f bpm) ---\n', Q_inj_bpm);
disp(T_inj);
fprintf('Mass-balance closure: sum(Net_bpm) = %.3f bpm (target %.3f)\n', ...
    sum(T_inj.Net_bpm, 'omitnan'), -Q_inj_bpm);

% Injection apparent retardation over the reference interval above Stage 5.
mask_ref_inj = z_iso_inj>=REF_Z_MIN & z_iso_inj<=REF_Z_MAX & isfinite(z_iso_inj) & ...
               moving_down & ~mask_slope_inj;
if sum(mask_ref_inj) >= 3
    p_ref_inj     = polyfit(t_min(mask_ref_inj), z_iso_inj(mask_ref_inj), 1);
    v_iso_ref_inj = p_ref_inj(1);              % ft/min, positive = down
    R_inj         = v_piston_inj / v_iso_ref_inj;
else
    v_iso_ref_inj = NaN;  R_inj = NaN;  p_ref_inj = [NaN NaN];
    warning('Insufficient points for injection reference-interval fit (%d)', sum(mask_ref_inj));
end

%% =========================================================
%           PRODUCTION isotherm slope analysis
% =========================================================
z_iso_prod = contour_front_intersections( ...
    t_min, z_dts, T_test, T_iso, [Z_TRACK_MIN Z_TRACK_MAX], [PROD_T_MIN PROD_T_MAX], -1, ...
    CONTOUR_MAX_FILL_GAP_STEPS);
z_iso_prod = z_iso_prod(:);
moving_up  = isfinite(z_iso_prod) & t_min(:)>=PROD_T_MIN & t_min(:)<=PROD_T_MAX;

[tf_prod, zf_prod] = build_fit_points(t_min, z_iso_prod, moving_up);
mask_slopefit_prod = points_in_depth_mask(zf_prod, MASK_Z_PROD);
tf_prod = tf_prod(~mask_slopefit_prod);  zf_prod = zf_prod(~mask_slopefit_prod);
mask_slope_prod = points_in_depth_mask(z_iso_prod, MASK_Z_PROD);

nP = numel(fit_stages_prod);
v_deep_prod = nan(nP,1);  v_shallow_prod = nan(nP,1);
ratio_prod = nan(nP,1);    % v_shallow/v_deep; ratios <1 are retained as losses
z_cross_prod = nan(nP,1); t_cross_prod = nan(nP,1);
p_deep_prod_all = nan(nP,2);  p_shallow_prod_all = nan(nP,2);   % stored for plotting

fprintf('\n=== PRODUCTION isotherm slope analysis (T_iso=%.0f C, Qprod=%.2f bpm) ===\n', T_iso, Q_prod_bpm);
fprintf('Piston velocity (annulus): %.2f ft/min\n', v_piston_prod);
if ~any(strcmpi(fit_stages_prod, 'Stage3'))
    fprintf('Stage3 omitted: no resolvable production-phase slope break.\n');
end
fprintf('%-10s %10s %10s %10s %10s\n', 'Stage','v_deep','v_shallow','ratio','Qchg');
fprintf('%s\n', repmat('-',1,55));
for si = 1:nP
    lims = PROD_FIT.(fit_stages_prod{si});
    [deep_min, deep_max, shallow_min, shallow_max] = unpack_prod_fit(lims);
    mask_deep    = zf_prod>=deep_min    & zf_prod<=deep_max;
    mask_shallow = zf_prod>=shallow_min & zf_prod<=shallow_max;
    if sum(mask_deep)<3 || sum(mask_shallow)<3
        fprintf('%-10s  insufficient points (deep=%d, shallow=%d)\n', ...
            fit_stages_prod{si}, sum(mask_deep), sum(mask_shallow)); continue
    end
    p_deep    = polyfit(tf_prod(mask_deep),    zf_prod(mask_deep),    1);
    p_shallow = polyfit(tf_prod(mask_shallow), zf_prod(mask_shallow), 1);
    p_deep_prod_all(si,:) = p_deep;  p_shallow_prod_all(si,:) = p_shallow;
    v_deep_prod(si)    = abs(p_deep(1));
    v_shallow_prod(si) = abs(p_shallow(1));
    ratio_prod(si)     = v_shallow_prod(si) / v_deep_prod(si);  % >1 gain, <1 loss upward
    t_cross_prod(si)   = (p_deep(2) - p_shallow(2)) / (p_shallow(1) - p_deep(1));
    z_cross_prod(si)   = p_deep(1)*t_cross_prod(si) + p_deep(2);   % breakpoint depth
    fprintf('%-10s %10.2f %10.2f %10.3f %10.3f\n', ...
        fit_stages_prod{si}, v_deep_prod(si), v_shallow_prod(si), ratio_prod(si), ratio_prod(si)-1);
end

% Production flow balance, anchored at the metered surface rate and walked
% shallow -> deep, so the open-hole (barefoot) inflow falls out as a derived
% result rather than being prescribed.
fit_stages_prod_topdown = fliplr(fit_stages_prod);   % Stage5,4,2,1
ratio_prod_topdown      = flipud(ratio_prod);
Qt = Q_prod_bpm;
Q_track_prod_topdown = nan(nP,1);   % flow below each stage, shallow->deep order
for si = 1:nP
    if isnan(ratio_prod_topdown(si)), continue; end
    Qt = Qt / ratio_prod_topdown(si);            % Q_deep = Q_shallow / (v_shallow/v_deep)
    Q_track_prod_topdown(si) = Qt;
end
Q_track_prod = flipud(Q_track_prod_topdown);   % back to deep->shallow order

Q_below_prod_full = fill_passthrough(stage_order_sd, fit_stages_prod_topdown, Q_track_prod_topdown, Q_prod_bpm);
T_prod = build_flow_balance(stage_order_sd, Q_below_prod_full, Q_prod_bpm, 'produce', 'Barefoot');
fprintf('\n--- PRODUCTION flow balance (anchored at surface Qprod=%.3f bpm) ---\n', Q_prod_bpm);
disp(T_prod);
fprintf('Mass-balance closure: sum(Net_bpm) = %.3f bpm (target %.3f)\n', ...
    sum(T_prod.Net_bpm, 'omitnan'), Q_prod_bpm);

% Production apparent retardation over the reference interval above Stage 5.
mask_ref_prod = z_iso_prod>=REF_Z_MIN & z_iso_prod<=REF_Z_MAX & isfinite(z_iso_prod) & ...
                moving_up & ~mask_slope_prod;
if sum(mask_ref_prod) >= 3
    p_ref_prod     = polyfit(t_min(mask_ref_prod), z_iso_prod(mask_ref_prod), 1);
    v_iso_ref_prod = abs(p_ref_prod(1));       % ft/min, abs because the front moves up
    R_prod         = v_piston_prod / v_iso_ref_prod;
else
    v_iso_ref_prod = NaN;  R_prod = NaN;  p_ref_prod = [NaN NaN];
    warning('Insufficient points for production reference-interval fit (%d)', sum(mask_ref_prod));
end

%% =========================================================
%      APPARENT retardation above Stage 5 (Section 3.3)
% =========================================================
% R = v_piston / v_iso_ref, from the metered rate and the annulus area only.
% This is the observed thermal-front retardation in a no-exchange interval. It
% is not predicted from casing/CT heat capacity here, and it does NOT enter the
% flow estimate, which uses velocity ratios: cross-section and retardation
% cancel in the ratio.
fprintf('\n=== APPARENT retardation, reference interval %.0f-%.0f ft (above Stage 5) ===\n', ...
    REF_Z_MIN, REF_Z_MAX);
fprintf('%-12s %12s %12s %12s\n', 'Phase','v_piston','v_iso_ref','R=vp/vi');
fprintf('%s\n', repmat('-',1,52));
fprintf('%-12s %12.2f %12.2f %12.3f\n', 'Injection',  v_piston_inj,  v_iso_ref_inj,  R_inj);
fprintf('%-12s %12.2f %12.2f %12.3f\n', 'Production', v_piston_prod, v_iso_ref_prod, R_prod);

%% =========================================================
%      DERIVED FLOW PROFILE
% =========================================================
% The breakpoint depth is the intersection of the shallow and deep fit lines
% (z_cross), not the stage midpoint. The flow fraction steps at each
% breakpoint to build Q(z)/Q_ref.
z_prof = z_dts;

% Injection: 100% above the shallowest fitted stage, multiplied by
% (1 - Qf/Qu) at each breakpoint.
Q_deep_stage_inj = nan(nS,1);
Q_track = 1.0;
for si = 1:nS
    if ~isnan(Qf_Qu_inj(si)), Q_track = Q_track * (1 - Qf_Qu_inj(si)); end
    Q_deep_stage_inj(si) = Q_track;
end
q_step_inj = ones(numel(z_prof),1);
for si = 1:nS
    if isnan(z_cross_inj(si)), continue; end
    [~, k] = min(abs(z_prof - z_cross_inj(si)));
    q_step_inj(k:end) = Q_deep_stage_inj(si);
end

% Production: normalized to 100% above Stage 5, walking shallow -> deep and
% dividing by ratio_prod (= v_shallow/v_deep, so Q_deep = Q_shallow/ratio).
% This is the local axial wellbore flow normalized by the metered surface rate,
% which is what Figure 5 plots.
fit_stages_prod_rev = fliplr(fit_stages_prod);   % shallow->deep
q_step_prod = ones(numel(z_prof),1);
Q_track2 = 1.0;  prev_k = 1;
for si = 1:numel(fit_stages_prod_rev)
    idx = find(strcmp(fit_stages_prod, fit_stages_prod_rev{si}));
    if isnan(z_cross_prod(idx)), continue; end
    [~, k] = min(abs(z_prof - z_cross_prod(idx)));
    q_step_prod(prev_k:k) = Q_track2;
    if ~isnan(ratio_prod(idx)), Q_track2 = Q_track2 / ratio_prod(idx); end
    prev_k = k + 1;
end
q_step_prod(prev_k:end) = Q_track2;

%% =========================================================
%                    COLLECTED RESULTS
% =========================================================
RESULTS = struct();
RESULTS.T_iso                 = T_iso;
RESULTS.injection.stages      = fit_stages_inj;
RESULTS.injection.Qf_frac     = Qf_Qu_inj;     % fraction exchanged at each stage
RESULTS.injection.Ql_frac     = Ql_Qu_inj;     % v_deep/v_shallow
RESULTS.injection.Qf_bpm      = Qf_bpm_inj;    % absolute loss at each stage, bpm
RESULTS.injection.z_cross     = z_cross_inj;   % breakpoint depth, ft
RESULTS.injection.R_apparent  = R_inj;
RESULTS.production.stages     = fit_stages_prod;
RESULTS.production.ratio      = ratio_prod;    % v_shallow/v_deep
RESULTS.production.Q_bpm      = Q_track_prod;  % flow below each stage, deep->shallow
RESULTS.production.z_cross    = z_cross_prod;  % breakpoint depth, ft
RESULTS.production.R_apparent = R_prod;
RESULTS.profile.z             = z_prof;        % depth, ft
RESULTS.profile.q_inj_frac    = q_step_inj;    % Q(z)/Qinj
RESULTS.profile.q_prod_frac   = q_step_prod;   % Q(z)/Qprod, local axial
RESULTS.tracking.z_inj        = z_iso_inj;
RESULTS.tracking.z_prod       = z_iso_prod;
RESULTS.tracking.mask_slope_inj  = mask_slope_inj;
RESULTS.tracking.mask_slope_prod = mask_slope_prod;

% Per-stage breakpoint summary.
stg_all = fieldnames(FRAC_STAGES);
nA = numel(stg_all);
PerfTop = nan(nA,1);  PerfBot = nan(nA,1);
Zcross_inj = nan(nA,1);  Zcross_prod = nan(nA,1);
for k = 1:nA
    [PerfTop(k), PerfBot(k)] = stage_depth_extent(FRAC_STAGES, stg_all{k});
    ii = find(strcmp(fit_stages_inj,  stg_all{k}), 1);
    jj = find(strcmp(fit_stages_prod, stg_all{k}), 1);
    if ~isempty(ii), Zcross_inj(k)  = z_cross_inj(ii);  end
    if ~isempty(jj), Zcross_prod(k) = z_cross_prod(jj); end
end
T_bp = table(string(stg_all(:)), PerfTop, PerfBot, round(Zcross_inj,1), round(Zcross_prod,1), ...
    'VariableNames', {'Stage','Perf_top_ft','Perf_bot_ft','Zcross_inj_ft','Zcross_prod_ft'});
RESULTS.tables = struct('injection', T_inj, 'production', T_prod, 'breakpoints', T_bp);

fprintf('\n--- Stage breakpoints (ft MD) ---\n');
disp(T_bp);

%% =========================================================
%                        PLOTTING
% =========================================================
stage_names_all = fieldnames(FRAC_STAGES);
stage_order     = get_stage_order_deep_to_shallow(FRAC_STAGES, stage_names_all);
colors_stage    = lines(numel(stage_order));
contour_levels  = 150:5:220;
t_prod_min_val  = minutes(PRODUCTION_TIME - INJECTION_TIME);
t_shut_min_val  = minutes(SHUTIN_TIME     - INJECTION_TIME);

% Depth axis shown in metres. The data stay in feet; only the window and the
% tick labels change, so the physical window is identical either way.
DEPTH_AXIS_M = [2600 3070];   % plotted depth window, m (2600 m = 8530 ft, 3070 m = 10073 ft)
DEPTH_TICK_M = 100;           % metre tick spacing
DEPTH_UNITS  = 'm';           % 'm' for the published figures; 'ft' for native working ticks
depth_ylabel = 'Depth MD (ft KB)'; if strcmpi(DEPTH_UNITS,'m'), depth_ylabel = 'Depth MD (m)'; end

% -----------------------------------------------------
% Figure 2: DTS waterfall, contour lines colored by temperature (no fill).
% clim is set to the contour range so line color maps directly to temperature.
% The tracked isotherm is overplotted in heavy black.
% -----------------------------------------------------
figure('Name','Figure 2 - DTS waterfall, Test 1','Position',[100 100 1000 800]);
ax_wf = axes; hold(ax_wf,'on');
set(ax_wf,'YDir','reverse','FontSize',14);
colormap(ax_wf, jet);
[Cwf, hwf] = contour(ax_wf, t_min, z_dts, T_test', contour_levels, 'LineWidth', 1.1);
[Ciso, hiso] = contour(ax_wf, t_min, z_dts, T_test', [T_iso T_iso], 'LineColor','k', 'LineWidth', 2.2);
clabel(Ciso, hiso, 'FontSize', 12, 'FontWeight','bold', 'LabelSpacing', 400, 'Color', 'k');
clim(ax_wf, [min(contour_levels) max(contour_levels)]);
clabel(Cwf, hwf, setdiff(contour_levels, T_iso), 'FontSize', 10, 'LabelSpacing', 400, 'Color', 'k');
cb = colorbar(ax_wf);
cb.Label.String = sprintf('Temperature (%s)', DTS_TEMP_UNITS);
cb.Label.FontSize = 14;
set_depth_axis(ax_wf, DEPTH_AXIS_M(1), DEPTH_AXIS_M(2), DEPTH_TICK_M, DEPTH_UNITS);
xlim(ax_wf, [0 400]);
xlabel(ax_wf, 'Time from injection start (min)', 'FontSize', 16);
ylabel(ax_wf, depth_ylabel, 'FontSize', 16);
grid(ax_wf, 'on');
add_waterfall_annotations(ax_wf, stage_order, FRAC_STAGES, colors_stage, ...
    MASK_Z_INJ, MASK_Z_PROD, INJ_T_MAX, PROD_T_MIN, PROD_T_MAX, ...
    t_prod_min_val, t_shut_min_val);
box(ax_wf, 'on'); hold(ax_wf, 'off');

% -----------------------------------------------------
% Figure 5: injection tracking | production tracking | derived flow profile,
% on a shared depth axis so stages align across panels.
% -----------------------------------------------------
% The injection duration sets the time span for both tracking panels, so a
% given ft/min draws at the same angle in each: panel 1 spans
% [INJ_T_MIN, INJ_T_MAX] and panel 2 spans [PROD_T_MIN, PROD_T_MIN + that same
% duration]. Production tail beyond that span is off-panel.
t_span = INJ_T_MAX - INJ_T_MIN;

stage_color = @(sn) colors_stage(strcmp(stage_order, sn), :);   % stage name -> boundary color
figure('Name','Figure 5 - Isotherm tracking and flow profile, Test 1','Position',[50 250 1250 900]);

% --- Panel 1: injection tracking ---
ax_inj = subplot(1,3,1); hold(ax_inj,'on');
if any(mask_slope_inj)
    plot(ax_inj, t_min(mask_slope_inj), z_iso_inj(mask_slope_inj), '.', ...
        'Color', [0.50 0.50 0.50], 'MarkerSize',10, ...
        'DisplayName', 'omitted (local circulation)');
end
plot(ax_inj, t_min(~mask_slope_inj & isfinite(z_iso_inj)), z_iso_inj(~mask_slope_inj & isfinite(z_iso_inj)), ...
    'k.', 'MarkerSize',6, 'DisplayName', sprintf('T=%.0f\\circC tracked front', T_iso));
kf = find(isfinite(z_iso_inj),1,'first');
t_pist_vec = t_min(kf):1:INJ_T_MAX;
z_pist_vec = z_iso_inj(kf) + v_piston_inj*(t_pist_vec - t_min(kf));
plot(ax_inj, t_pist_vec, z_pist_vec, 'g--', 'LineWidth',1.5, ...
    'DisplayName', sprintf('Piston: v=%.1f m/min', v_piston_inj*FTMIN_TO_MMIN));
if ~isnan(v_iso_ref_inj)
    plot(ax_inj, t_min(mask_ref_inj), polyval(p_ref_inj, t_min(mask_ref_inj)), 'm-', 'LineWidth',2, ...
        'DisplayName', sprintf('Ref interval: v=%.1f m/min (R=%.2f)', v_iso_ref_inj*FTMIN_TO_MMIN, R_inj));
end
for si = 1:nS
    if all(isnan(p_shallow_inj_all(si,:))), continue; end
    sn = fit_stages_inj{si};  cs = stage_color(sn);
    [shallow_min, shallow_max, deep_min, deep_max] = unpack_inj_fit(INJ_FIT.(sn));
    m_s = z_iso_inj>=shallow_min & z_iso_inj<=shallow_max & isfinite(z_iso_inj) & moving_down & ~mask_slope_inj;
    m_d = z_iso_inj>=deep_min    & z_iso_inj<=deep_max    & isfinite(z_iso_inj) & moving_down & ~mask_slope_inj;
    % Solid = fit above the stage, dashed = fit below the stage.
    plot(ax_inj, t_min(m_s), polyval(p_shallow_inj_all(si,:), t_min(m_s)), '-',  'Color', cs, 'LineWidth',2, 'HandleVisibility','off');
    plot(ax_inj, t_min(m_d), polyval(p_deep_inj_all(si,:),    t_min(m_d)), '--', 'Color', cs, 'LineWidth',2, 'HandleVisibility','off');
    if ~isnan(z_cross_inj(si))
        plot(ax_inj, t_cross_inj(si), z_cross_inj(si), 'x', 'Color', cs, ...
            'MarkerSize',12, 'LineWidth',2, 'HandleVisibility','off');
    end
end
if any(mask_slope_inj)
    plot(ax_inj, t_min(mask_slope_inj), z_iso_inj(mask_slope_inj), '.', ...
        'Color', [0.50 0.50 0.50], 'MarkerSize',10, 'HandleVisibility','off');
end
draw_stage_lines(ax_inj, stage_order, FRAC_STAGES, colors_stage, 5);
set_depth_axis(ax_inj, DEPTH_AXIS_M(1), DEPTH_AXIS_M(2), DEPTH_TICK_M, DEPTH_UNITS);
xlim(ax_inj, [INJ_T_MIN, INJ_T_MIN + t_span]);
set(ax_inj, 'YDir','reverse', 'FontSize',10);
xlabel(ax_inj, 'Time from injection start (min)', 'FontSize',12);
ylabel(ax_inj, depth_ylabel, 'FontSize',12);
title(ax_inj, sprintf('Injection T=%.0f\\circC', T_iso), 'FontSize',12);
legend(ax_inj, 'Location','northeast', 'FontSize',10); grid(ax_inj,'on'); hold(ax_inj,'off');

% --- Panel 2: production tracking ---
ax_prod = subplot(1,3,2); hold(ax_prod,'on');
if any(mask_slope_prod)
    plot(ax_prod, t_min(mask_slope_prod), z_iso_prod(mask_slope_prod), '.', ...
        'Color', [0.50 0.50 0.50], 'MarkerSize',10, ...
        'DisplayName', 'omitted (local circulation)');
end
plot(ax_prod, t_min(~mask_slope_prod & isfinite(z_iso_prod)), z_iso_prod(~mask_slope_prod & isfinite(z_iso_prod)), ...
    'k.', 'MarkerSize',6, 'DisplayName', sprintf('T=%.0f\\circC tracked front', T_iso));
kfp = find(t_min>=PROD_T_MIN & isfinite(z_iso_prod), 1, 'first');
t_pist_vec_p = PROD_T_MIN:1:PROD_T_MAX;
z_pist_vec_p = z_iso_prod(kfp) - v_piston_prod*(t_pist_vec_p - PROD_T_MIN);
plot(ax_prod, t_pist_vec_p, z_pist_vec_p, 'g--', 'LineWidth',1.5, ...
    'DisplayName', sprintf('Piston: v=%.1f m/min', v_piston_prod*FTMIN_TO_MMIN));
if ~isnan(v_iso_ref_prod)
    plot(ax_prod, t_min(mask_ref_prod), polyval(p_ref_prod, t_min(mask_ref_prod)), 'm-', 'LineWidth',2, ...
        'DisplayName', sprintf('Ref interval: v=%.1f m/min (R=%.2f)', v_iso_ref_prod*FTMIN_TO_MMIN, R_prod));
end
for si = 1:nP
    if all(isnan(p_deep_prod_all(si,:))), continue; end
    sn = fit_stages_prod{si};  cs = stage_color(sn);
    [deep_min, deep_max, shallow_min, shallow_max] = unpack_prod_fit(PROD_FIT.(sn));
    m_d = z_iso_prod>=deep_min    & z_iso_prod<=deep_max    & isfinite(z_iso_prod) & moving_up & ~mask_slope_prod;
    m_s = z_iso_prod>=shallow_min & z_iso_prod<=shallow_max & isfinite(z_iso_prod) & moving_up & ~mask_slope_prod;
    plot(ax_prod, t_min(m_s), polyval(p_shallow_prod_all(si,:), t_min(m_s)), '-',  'Color', cs, 'LineWidth',2, 'HandleVisibility','off');
    plot(ax_prod, t_min(m_d), polyval(p_deep_prod_all(si,:),    t_min(m_d)), '--', 'Color', cs, 'LineWidth',2, 'HandleVisibility','off');
    if ~isnan(z_cross_prod(si))
        plot(ax_prod, t_cross_prod(si), z_cross_prod(si), 'x', 'Color', cs, ...
            'MarkerSize',14, 'LineWidth',2.5, 'HandleVisibility','off');
    end
end
if any(mask_slope_prod)
    plot(ax_prod, t_min(mask_slope_prod), z_iso_prod(mask_slope_prod), '.', ...
        'Color', [0.50 0.50 0.50], 'MarkerSize',10, 'HandleVisibility','off');
end
draw_stage_lines(ax_prod, stage_order, FRAC_STAGES, colors_stage, []);
set_depth_axis(ax_prod, DEPTH_AXIS_M(1), DEPTH_AXIS_M(2), DEPTH_TICK_M, DEPTH_UNITS);
xlim(ax_prod, [PROD_T_MIN, PROD_T_MIN + t_span]);
set(ax_prod, 'YDir','reverse', 'FontSize',10);
xlabel(ax_prod, 'Time from injection start (min)', 'FontSize',12);
ylabel(ax_prod, '');  set(ax_prod, 'YTickLabel', []);   % shared depth axis with panel 1
title(ax_prod, sprintf('Production T=%.0f\\circC', T_iso), 'FontSize',12);
legend(ax_prod, 'Location','southeast', 'FontSize',10); grid(ax_prod,'on'); hold(ax_prod,'off');

% --- Panel 3: derived flow profile ---
ax_q = subplot(1,3,3); hold(ax_q,'on');
plot(ax_q, q_step_inj*100,  z_prof, 'b-', 'LineWidth',2, 'DisplayName','Injection');
plot(ax_q, q_step_prod*100, z_prof, 'r-', 'LineWidth',2, 'DisplayName','Production');

% Each flow step is labelled at the midpoint of its vertical segment, halfway
% between its own breakpoint and the next deeper breakpoint, so the label sits
% centred on the constant-flow run. Label depths are in feet (data units).
z_win_bot_ft  = DEPTH_AXIS_M(2) / 0.3048;              % bottom of plotted window, ft
zc_inj_finite = z_cross_inj(isfinite(z_cross_inj));
for si = 1:numel(fit_stages_inj)
    if isnan(z_cross_inj(si)) || isnan(Q_deep_stage_inj(si)), continue; end
    deeper = zc_inj_finite(zc_inj_finite > z_cross_inj(si));
    z_bot  = z_win_bot_ft; if ~isempty(deeper), z_bot = min(deeper); end
    z_lbl  = 0.5*(z_cross_inj(si) + z_bot);
    text(ax_q, Q_deep_stage_inj(si)*100 - 1.5, z_lbl, ...
        sprintf('%.1f m^3/hr', Q_deep_stage_inj(si)*Q_inj_bpm*BPM_TO_M3HR), ...
        'Color','b', 'FontSize',10, 'HorizontalAlignment','right', ...
        'VerticalAlignment','middle', 'Clipping','on');
end

% Manual offsets for the production flow labels, keyed by stage, to stop the
% deep labels overlapping. dx is in % (panel x units, + = right); dy is in FEET
% (the axis reads metres but the data are feet; + = deeper).
PROD_LBL_DX = struct('Stage1',22,'Stage2',25,'Stage3',0,'Stage4',0,'Stage5',0);
PROD_LBL_DY = struct('Stage1',23,'Stage2',-2,'Stage3',0,'Stage4',0,'Stage5',0);
zc_prod_finite = z_cross_prod(isfinite(z_cross_prod));
for si = 1:numel(fit_stages_prod)
    if isnan(z_cross_prod(si)) || isnan(Q_track_prod(si)), continue; end
    sn = fit_stages_prod{si};
    deeper = zc_prod_finite(zc_prod_finite > z_cross_prod(si));
    z_bot  = z_win_bot_ft; if ~isempty(deeper), z_bot = min(deeper); end
    z_lbl  = 0.5*(z_cross_prod(si) + z_bot);
    text(ax_q, (Q_track_prod(si)/Q_prod_bpm)*100 - 1.5 + PROD_LBL_DX.(sn), ...
               z_lbl + PROD_LBL_DY.(sn), ...
        sprintf('%.1f m^3/hr', Q_track_prod(si)*BPM_TO_M3HR), ...
        'Color','r', 'FontSize',10, 'HorizontalAlignment','right', ...
        'VerticalAlignment','middle', 'Clipping','on');
end

% Surface reference anchors. These are not fitted stages, so they have no
% z_cross and are labelled explicitly at the top of the panel.
text(ax_q, 100 - 1.5, Z_TRACK_MIN + 30, sprintf('%.1f m^3/hr', Q_inj_bpm*BPM_TO_M3HR), ...
    'Color','b', 'FontSize',10, 'FontWeight','bold', 'HorizontalAlignment','right', ...
    'VerticalAlignment','top', 'Clipping','on');
text(ax_q, 100 - 1.5, Z_TRACK_MIN + 90, sprintf('%.1f m^3/hr', Q_prod_bpm*BPM_TO_M3HR), ...
    'Color','r', 'FontSize',10, 'FontWeight','bold', 'HorizontalAlignment','right', ...
    'VerticalAlignment','top', 'Clipping','on');

draw_stage_lines(ax_q, stage_order, FRAC_STAGES, colors_stage, []);
set(ax_q, 'YDir','reverse', 'FontSize',10);
set_depth_axis(ax_q, DEPTH_AXIS_M(1), DEPTH_AXIS_M(2), DEPTH_TICK_M, DEPTH_UNITS);
xlim(ax_q, [0 101]);
xlabel(ax_q, 'Flow profile / Q_{reference} (%)', 'FontSize',12);
ylabel(ax_q, '');  set(ax_q, 'YTickLabel', []);   % shared depth axis with panel 1
title(ax_q, 'Derived Flow Profile', 'FontSize',12);
legend(ax_q, 'Location','northwest', 'FontSize',10); grid(ax_q,'on'); hold(ax_q,'off');

ax_inj.Position  = [0.06 0.12 0.30 0.78];
ax_prod.Position = [0.38 0.12 0.30 0.78];
ax_q.Position    = [0.70 0.12 0.28 0.78];

fprintf('\nDone. Flow-partitioning outputs collected in struct RESULTS.\n');

% =========================================================
%                     HELPER FUNCTIONS
% =========================================================

function [shallow_min, shallow_max, deep_min, deep_max] = unpack_inj_fit(lims)
% Injection input convention:
%   [shallow_min shallow_max deep_min deep_max]
% The pairs are sorted internally so accidental reversed min/max order does not
% change the fit window.
    s = sort(lims(1:2));
    d = sort(lims(3:4));
    shallow_min = s(1); shallow_max = s(2);
    deep_min    = d(1); deep_max    = d(2);
end

function [deep_min, deep_max, shallow_min, shallow_max] = unpack_prod_fit(lims)
% Production input convention (same as INJ_FIT, top-down):
%   [shallow_min shallow_max deep_min deep_max]
% Pairs are sorted internally, so either order within a pair is tolerated.
    s = sort(lims(1:2));
    d = sort(lims(3:4));
    shallow_min = s(1);  shallow_max = s(2);
    deep_min    = d(1);  deep_max    = d(2);
end


function [zf, info, candByTime, selectedBranch] = contour_front_intersections( ...
    t_min, z_dts, T_test, T_iso, zlim, twin, direction, maxFillGapSteps)
% Extract a single-valued isotherm front from the 2-D contour field.
%   direction = +1 for injection: select deepest crossing at each time.
%   direction = -1 for production: select shallowest crossing at each time.
    t_min = double(t_min(:));
    z_dts = double(z_dts(:))';

    C = contourc(t_min', z_dts, double(T_test'), [T_iso T_iso]);
    branches = parse_contour_branches(C, zlim, twin);

    zf = nan(numel(t_min),1);
    candByTime = cell(numel(t_min),1);
    selectedBranch = [];
    info = struct();

    if isempty(branches)
        info.message = 'No branches found.';
        info.bestIdx = NaN;
        return
    end

    scores = nan(numel(branches),1);
    branchStats = nan(numel(branches),6);
    for b = 1:numel(branches)
        [scores(b), branchStats(b,:)] = score_raw_branch(branches{b}, direction);
    end

    [~, bestIdx] = max(scores);
    selectedBranch = branches{bestIdx};
    info.bestIdx = bestIdx;
    info.scores = scores;
    info.branchStats = branchStats;
    info.twin = twin;

    active = t_min >= twin(1) & t_min <= twin(2);
    for i = find(active(:))'
        candByTime{i} = vertical_polyline_intersections(selectedBranch, t_min(i), zlim);
    end

    for i = find(active(:))'
        zc = candByTime{i};
        if isempty(zc), continue; end
        zc = sort(zc(:));
        if direction > 0
            zf(i) = max(zc);  % injection: advancing/deep envelope
        else
            zf(i) = min(zc);  % production: returning/shallow envelope
        end
    end

    try
        zf = fillmissing(zf, 'linear', 'MaxGap', maxFillGapSteps);
    catch
        zf = fillmissing(zf, 'linear');
    end
    info.nFinite = nnz(isfinite(zf));
end

function branches = parse_contour_branches(C, zlim, twin)
    branches = {};
    k = 1;
    while k < size(C,2)
        n = C(2,k);
        tt = C(1,k+1:k+n)';
        zz = C(2,k+1:k+n)';
        keep = tt >= twin(1) & tt <= twin(2) & ...
               zz >= zlim(1) & zz <= zlim(2) & ...
               isfinite(tt) & isfinite(zz);
        tt = tt(keep); zz = zz(keep);
        if numel(tt) >= 3
            branches{end+1} = [tt zz]; %#ok<AGROW>
        end
        k = k + n + 1;
    end
end

function [score, stats] = score_raw_branch(s, direction)
    tt = s(:,1); zz = s(:,2);
    if numel(tt) < 3
        score = -Inf; stats = nan(1,6); return
    end
    dt = diff(tt); dz = diff(zz);
    ok = abs(dt) > eps & isfinite(dt) & isfinite(dz);
    if any(ok)
        localSlope = dz(ok) ./ dt(ok);
        dirOK = mean(sign(localSlope) == direction);
    else
        localSlope = NaN;
        dirOK = 0;
    end
    tspan = max(tt) - min(tt);
    zspan = max(zz) - min(zz);
    p = polyfit(tt, zz, 1);
    globalSlope = p(1);
    globalOK = max(0, sign(globalSlope) * direction);
    if numel(localSlope) >= 3
        roughness = median(abs(diff(localSlope)), 'omitnan');
    else
        roughness = 0;
    end
    score = 5*tspan + 1*zspan + 150*dirOK + 50*globalOK - 0.05*roughness;
    stats = [score, tspan, zspan, globalSlope, dirOK, roughness];
end

function zc = vertical_polyline_intersections(branch, ti, zlim)
    tt = branch(:,1); zz = branch(:,2);
    zc = [];
    for k = 1:numel(tt)-1
        t1 = tt(k); t2 = tt(k+1);
        z1 = zz(k); z2 = zz(k+1);
        if ~isfinite(t1) || ~isfinite(t2) || ~isfinite(z1) || ~isfinite(z2)
            continue
        end
        if abs(t2 - t1) < eps
            if abs(ti - t1) < 1e-9
                zc = [zc; z1; z2]; %#ok<AGROW>
            end
            continue
        end
        if ti >= min(t1,t2) && ti <= max(t1,t2)
            a = (ti - t1) / (t2 - t1);
            if a >= -1e-12 && a <= 1+1e-12
                zi = z1 + a*(z2 - z1);
                if zi >= zlim(1) && zi <= zlim(2)
                    zc = [zc; zi]; %#ok<AGROW>
                end
            end
        end
    end
    if isempty(zc), return; end
    zc = sort(zc);
    keep = [true; abs(diff(zc)) > 1.0];
    zc = zc(keep);
end


function [tf, zf] = build_fit_points(t_min, z_track, moving)
% Flat (t,z) point set for the slope fits: the single-valued tracked front,
% restricted to the phase's moving window. These are exactly the points drawn
% on the tracking panels, so a fitted line always lies on its own data.
    t_min = t_min(:);  z_track = z_track(:);  moving = logical(moving(:));
    keep = moving & isfinite(z_track);
    tf = t_min(keep);  zf = z_track(keep);
end

function mask = points_in_depth_mask(z_front, mask_z)
% Logical mask for selected-front points whose depth falls inside any [zlo zhi]
% interval in mask_z. Used only to exclude points from slope fits; tracking
% itself remains unmasked.
    z_front = z_front(:);
    mask = false(size(z_front));
    if isempty(mask_z), return; end
    for mi = 1:size(mask_z,1)
        zlo = min(mask_z(mi,:));
        zhi = max(mask_z(mi,:));
        mask = mask | (z_front >= zlo & z_front <= zhi);
    end
end

function qfull = fill_passthrough(order, fitStages, Q_below_fit, Q_ref)
% Map fitted per-stage flow onto the FULL stage list (order = shallow->deep).
% Stages that were not fit (or failed) carry the flow through unchanged from
% the stage above, so every stage gets a row and its net change reads as 0.
    n = numel(order);  qfull = nan(n,1);  prev = Q_ref;
    for k = 1:n
        idx = find(strcmp(fitStages, order{k}), 1);
        if ~isempty(idx) && isfinite(Q_below_fit(idx))
            qfull(k) = Q_below_fit(idx);
        else
            qfull(k) = prev;            % pass-through: no detected flow change here
        end
        prev = qfull(k);
    end
end

function T = build_flow_balance(stageNames, Q_below, Q_ref, mode, terminalLabel)
% Cumulative wellbore flow balance, anchored at the surface (metered) rate.
%   stageNames    : cell/string, ordered SURFACE -> DEEP
%   Q_below       : flow (bpm) remaining below each stage, same order
%   Q_ref         : surface rate, bpm
%   mode          : 'inject' (Net<0 = loss to formation) or
%                   'produce' (Net>0 = gain / inflow, moving up)
%   terminalLabel : terminal boundary below the deepest tracked stage, normally
%                   'Barefoot'. The terminal row closes the mass balance:
%                   injection loses the residual flow to barefoot, while
%                   production gains the residual flow from barefoot.
    stageNames = stageNames(:);  Q_below = Q_below(:);
    fin = isfinite(Q_below);
    sn  = string(stageNames(fin));  qb = Q_below(fin);
    if isempty(qb), T = table(); return; end

    qa = [Q_ref; qb(1:end-1)];                 % flow above each stage
    switch lower(mode)
        case 'produce'
            net = qa - qb;                     % + = gain (fluid enters, moving up)
            terminal_net = qb(end);            % gain from barefoot/open-hole interval
        case 'inject'
            net = qb - qa;                     % - = loss (fluid exits, moving down)
            terminal_net = -qb(end);           % residual loss to barefoot/open-hole interval
        otherwise
            error('mode must be ''inject'' or ''produce''');
    end

    Stage   = [ "Surface"; sn; string(terminalLabel) ];
    Q_bpm   = round([Q_ref; qb; qb(end)], 3);
    Pct_of_surface = round(100 * Q_bpm / Q_ref, 1);
    Net_bpm = round([NaN; net; terminal_net], 3);
    Effect  = strings(numel(Stage), 1);
    Effect(Net_bpm > 0)  = "gain";
    Effect(Net_bpm < 0)  = "loss";
    Effect(Net_bpm == 0) = "none";   % stage carried through with no detected flow change
    T = table(Stage, Q_bpm, Pct_of_surface, Net_bpm, Effect);
end
function [gaps, dt_nom, thr] = detect_time_gaps(t, gap_factor, min_gap_minutes)
    dt = minutes(diff(t)); dt_nom = median(dt);
    thr = max(gap_factor * dt_nom, min_gap_minutes);
    idx = find(dt > thr); gaps = cell(numel(idx), 2);
    for k = 1:numel(idx), gaps{k,1} = t(idx(k)); gaps{k,2} = t(idx(k)+1); end
end

function mask = mask_rows_from_gap_intervals(t, gap_intervals)
    mask = false(size(t));
    for k = 1:size(gap_intervals,1)
        mask = mask | (t >= gap_intervals{k,1} & t <= gap_intervals{k,2});
    end
end

function [top, bottom] = stage_depth_extent(FRAC_STAGES, stage_name)
    intervals = FRAC_STAGES.(stage_name);
    top = min(intervals(:)); bottom = max(intervals(:));
end

function stage_order = get_stage_order_deep_to_shallow(FRAC_STAGES, stage_names)
    depths = zeros(numel(stage_names),1);
    for k = 1:numel(stage_names), depths(k) = min(FRAC_STAGES.(stage_names{k})(:)); end
    [~, idx] = sort(depths, 'descend'); stage_order = stage_names(idx);
end

function draw_stage_lines(ax, stage_order, FRAC_STAGES, colors_stage, label_x)
% Per-stage boundary lines (top/bottom of each stage's perf interval),
% colored by stage. If label_x is nonempty, also draw the stage name at
% that x-position (left panel only). Matches the synthetic figure: stage
% top/bottom only, no per-perforation sub-lines.
    for si = 1:numel(stage_order)
        sn = stage_order{si}; [st_top, st_bot] = stage_depth_extent(FRAC_STAGES, sn);
        yline(ax, st_top, '--', 'Color', colors_stage(si,:), 'LineWidth',1.2, 'Alpha',0.9, 'HandleVisibility','off');
        yline(ax, st_bot, '--', 'Color', colors_stage(si,:), 'LineWidth',1.2, 'Alpha',0.9, 'HandleVisibility','off');
        if ~isempty(label_x)
            text(ax, label_x, 0.5*(st_top+st_bot), sn, 'Color', colors_stage(si,:), ...
                'FontSize',11, 'FontWeight','bold', 'VerticalAlignment','middle', 'HorizontalAlignment','left');
        end
    end
end

function set_depth_axis(ax, z_lo_m, z_hi_m, dz_m, units)
% Depth axis for the plots. The DATA is always in feet; this only sets the
% window and the tick labels. The physical window is identical either way
% (the metre range converted to feet), so switching units never moves data.
%   units='ft': native feet ticks - best for working, since the label
%               nudges (PROD_LBL_DY, z_cross_*) are all in feet.
%   units='m' : ticks relabelled to round metres for the publication figure.
% On shared-axis panels the caller may blank labels afterwards with
% set(ax,'YTickLabel',[]); the tick POSITIONS still line up across panels.
    M2FT = 1/0.3048;
    ylim(ax, [z_lo_m z_hi_m]*M2FT);
    if strcmpi(units, 'm')
        m0 = ceil(z_lo_m/dz_m)*dz_m;  m1 = floor(z_hi_m/dz_m)*dz_m;
        m_ticks = m0:dz_m:m1;
        set(ax, 'YTick', m_ticks*M2FT, 'YTickLabel', compose('%d', m_ticks));
    else
        set(ax, 'YTickMode','auto', 'YTickLabelMode','auto');   % native feet ticks
    end
end

function add_waterfall_annotations(ax, stage_order, FRAC_STAGES, colors_stage, ...
    MASK_Z_INJ, MASK_Z_PROD, INJ_T_MAX, PROD_T_MIN, PROD_T_MAX, ...
    t_prod_min_val, t_shut_min_val)
% Slope-fit mask shaded intervals, per-stage boundary lines + labels, and
% injection/production/shut-in event lines for the waterfall figure.

    % Shaded boxes marking depth intervals excluded from slope fitting only.
    for mi = 1:size(MASK_Z_INJ, 1)
        patch(ax, [0 INJ_T_MAX INJ_T_MAX 0], ...
            [MASK_Z_INJ(mi,1) MASK_Z_INJ(mi,1) MASK_Z_INJ(mi,2) MASK_Z_INJ(mi,2)], ...
            [0.5 0.5 0.5], 'FaceAlpha', 0.25, 'EdgeColor', 'none', 'HandleVisibility', 'off');
    end
    for mi = 1:size(MASK_Z_PROD, 1)
        patch(ax, [PROD_T_MIN PROD_T_MAX PROD_T_MAX PROD_T_MIN], ...
            [MASK_Z_PROD(mi,1) MASK_Z_PROD(mi,1) MASK_Z_PROD(mi,2) MASK_Z_PROD(mi,2)], ...
            [0.5 0.5 0.5], 'FaceAlpha', 0.25, 'EdgeColor', 'none', 'HandleVisibility', 'off');
    end

    for si = 1:numel(stage_order)
        sn = stage_order{si}; [st_top, st_bot] = stage_depth_extent(FRAC_STAGES, sn);
        yline(ax, st_top, '--', 'Color', colors_stage(si,:), 'LineWidth', 1.5, 'Alpha', 0.9);
        yline(ax, st_bot, '--', 'Color', colors_stage(si,:), 'LineWidth', 1.5, 'Alpha', 0.9);
        text(ax, 2, 0.5*(st_top+st_bot), sn, 'Color', colors_stage(si,:), ...
            'FontSize', 13, 'FontWeight', 'bold', 'VerticalAlignment', 'middle');
    end
    xline(ax, t_prod_min_val, '-', 'Color', [0 0.7 0], 'LineWidth', 1.8, ...
        'Label', 'Production', 'LabelVerticalAlignment', 'top', 'LabelHorizontalAlignment', 'center', 'FontSize', 13);
    xline(ax, t_shut_min_val, '-', 'Color', [0.8 0 0], 'LineWidth', 1.8, ...
        'Label', 'Shut-in', 'LabelVerticalAlignment', 'top', 'LabelHorizontalAlignment', 'center', 'FontSize', 13);
end
