# Standard library imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import seaborn as sns
import warnings

# Scikit-learn: model, preprocessing, and evaluation
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error

warnings.filterwarnings('ignore')
np.random.seed(42)  # Reproducibility

# Aesthetic configuration for all plots
sns.set_theme(style='whitegrid', palette='muted', font_scale=1.1)
plt.rcParams.update({'figure.dpi': 120, 'axes.titlesize': 13})

print("Environment configured successfully.")

def generate_ml_signal(
    n_samples: int = 2000,
    noise_scale: float = 0.15,
    random_state: int = 42
) -> pd.DataFrame:
    """
    Generate a synthetic multivariate ML signal dataset with realistic properties.

    The target variable y is a nonlinear combination of lagged features,
    a cyclic regime feature, and heteroscedastic noise.

    Parameters
    ----------
    n_samples : int
        Total number of observations to generate.
    noise_scale : float
        Base scale of additive noise; actual noise is heteroscedastic.
    random_state : int
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with feature columns [f1..f6] and target column [y].
    """
    rng = np.random.default_rng(random_state)
    t = np.linspace(0, 4 * np.pi, n_samples)  # Time index for cyclic features

    # --- Feature Construction ---
    # f1: Momentum signal (autoregressive-like, mean zero)
    f1 = np.cumsum(rng.normal(0, 0.05, n_samples))
    f1 = (f1 - f1.mean()) / f1.std()           # Standardize to unit variance

    # f2: Mean-reversion signal (Ornstein-Uhlenbeck process approximation)
    f2 = np.zeros(n_samples)
    for i in range(1, n_samples):
        f2[i] = 0.95 * f2[i - 1] + rng.normal(0, 0.1)  # AR(1) with decay

    # f3: Cyclic regime indicator (captures periodic market regimes)
    f3 = np.sin(t) + 0.5 * np.cos(2 * t)

    # f4: Volatility proxy (squared returns of f1)
    f4 = np.abs(np.diff(f1, prepend=f1[0]))

    # f5: Cross-sectional rank signal (uniform noise, simulates ranking)
    f5 = rng.uniform(-1, 1, n_samples)

    # f6: Interaction term between momentum and regime
    f6 = f1 * f3

    # --- Target Construction ---
    # Nonlinear combination with interaction effects
    signal = (
        0.4 * f1
        - 0.3 * f2
        + 0.2 * np.tanh(f3)          # Nonlinear regime component
        + 0.15 * f4
        + 0.1 * f5
        + 0.25 * f6                  # Interaction feature contribution
    )

    # Heteroscedastic noise: variance scales with |signal|
    het_noise = noise_scale * (1 + 0.8 * np.abs(signal)) * rng.normal(0, 1, n_samples)
    y = signal + het_noise

    # Assemble into a DataFrame
    df = pd.DataFrame({
        'f1': f1, 'f2': f2, 'f3': f3,
        'f4': f4, 'f5': f5, 'f6': f6,
        'y': y
    })
    return df


# Generate the dataset
df = generate_ml_signal(n_samples=2000, noise_scale=0.15)

print(f"Dataset shape  : {df.shape}")
print(f"Target range   : [{df['y'].min():.3f}, {df['y'].max():.3f}]")
print(f"Target std dev : {df['y'].std():.4f}")
print("\nFirst 5 rows:")
df.head()

def split_dataset(
    df: pd.DataFrame,
    feature_cols: list,
    target_col: str = 'y',
    train_frac: float = 0.60,
    cal_frac: float = 0.20
) -> dict:
    """
    Partition a DataFrame into train, calibration, and test splits.

    Temporal ordering is preserved (no shuffle) to maintain exchangeability
    under the conformal prediction framework for sequential data.

    Parameters
    ----------
    df : pd.DataFrame
        Input DataFrame containing features and target.
    feature_cols : list of str
        Column names to use as predictive features.
    target_col : str
        Column name of the regression target.
    train_frac : float
        Proportion of data allocated to training.
    cal_frac : float
        Proportion of data allocated to calibration.

    Returns
    -------
    dict
        Dictionary with keys: X_train, X_cal, X_test, y_train, y_cal, y_test.
    """
    n = len(df)

    # Compute integer split indices
    train_end = int(n * train_frac)
    cal_end   = int(n * (train_frac + cal_frac))

    X = df[feature_cols].values   # Feature matrix as NumPy array
    y = df[target_col].values     # Target vector

    # Slice each partition sequentially
    splits = {
        'X_train': X[:train_end],
        'y_train': y[:train_end],
        'X_cal'  : X[train_end:cal_end],
        'y_cal'  : y[train_end:cal_end],
        'X_test' : X[cal_end:],
        'y_test' : y[cal_end:]
    }
    return splits


# Define feature columns (all except the target)
FEATURE_COLS = ['f1', 'f2', 'f3', 'f4', 'f5', 'f6']

# Execute the split
splits = split_dataset(df, feature_cols=FEATURE_COLS)

# Report partition sizes
for name in ['X_train', 'X_cal', 'X_test']:
    print(f"{name:10s}: {splits[name].shape[0]:5d} samples")

def scale_features(splits: dict) -> tuple[dict, StandardScaler]:
    """
    Apply zero-mean, unit-variance scaling to feature matrices.

    The scaler is fitted on the training partition only. The same fitted
    scaler is applied to calibration and test partitions to prevent leakage.

    Parameters
    ----------
    splits : dict
        Dictionary produced by split_dataset().

    Returns
    -------
    tuple
        (scaled_splits dict, fitted StandardScaler instance)
    """
    scaler = StandardScaler()

    # Fit on training data only
    scaler.fit(splits['X_train'])

    # Transform all three partitions using the same scaler
    scaled = dict(splits)   # Shallow copy to preserve y arrays
    scaled['X_train'] = scaler.transform(splits['X_train'])
    scaled['X_cal']   = scaler.transform(splits['X_cal'])
    scaled['X_test']  = scaler.transform(splits['X_test'])

    return scaled, scaler


# Apply scaling
splits_scaled, scaler = scale_features(splits)

print("Feature scaling applied.")
print(f"Train feature mean (post-scale): {splits_scaled['X_train'].mean(axis=0).round(4)}")
print(f"Train feature std  (post-scale): {splits_scaled['X_train'].std(axis=0).round(4)}")

def train_base_regressor(
    X_train: np.ndarray,
    y_train: np.ndarray,
    n_estimators: int = 300,
    learning_rate: float = 0.05,
    max_depth: int = 4,
    subsample: float = 0.8
) -> GradientBoostingRegressor:
    """
    Fit a Gradient Boosting Regressor on the training partition.

    Parameters
    ----------
    X_train : np.ndarray, shape (n_train, n_features)
        Scaled training feature matrix.
    y_train : np.ndarray, shape (n_train,)
        Training target vector.
    n_estimators : int
        Number of boosting stages.
    learning_rate : float
        Shrinkage factor applied to each tree's contribution.
    max_depth : int
        Maximum depth of individual regression estimators.
    subsample : float
        Fraction of samples used for fitting each base learner (stochastic GB).

    Returns
    -------
    GradientBoostingRegressor
        Fitted model instance.
    """
    model = GradientBoostingRegressor(
        n_estimators=n_estimators,
        learning_rate=learning_rate,
        max_depth=max_depth,
        subsample=subsample,
        loss='squared_error',   # Minimise MSE during boosting
        random_state=42
    )
    model.fit(X_train, y_train)
    return model


# Train the base regressor
model = train_base_regressor(
    splits_scaled['X_train'],
    splits_scaled['y_train']
)

# Evaluate on train and test for diagnostic purposes
train_preds = model.predict(splits_scaled['X_train'])
test_preds  = model.predict(splits_scaled['X_test'])

train_mae = mean_absolute_error(splits_scaled['y_train'], train_preds)
test_mae  = mean_absolute_error(splits_scaled['y_test'],  test_preds)
test_rmse = np.sqrt(mean_squared_error(splits_scaled['y_test'], test_preds))

print(f"Base Regressor Performance")
print(f"  Train MAE : {train_mae:.4f}")
print(f"  Test  MAE : {test_mae:.4f}")
print(f"  Test  RMSE: {test_rmse:.4f}")

def compute_nonconformity_scores(
    model: GradientBoostingRegressor,
    X_cal: np.ndarray,
    y_cal: np.ndarray
) -> np.ndarray:
    """
    Compute absolute residual nonconformity scores on the calibration set.

    The nonconformity score for observation i is defined as:
        s_i = |y_i - ŷ_i|
    where ŷ_i is the base model's point prediction.

    Parameters
    ----------
    model : fitted regressor
        Any fitted sklearn-compatible regressor with a predict() method.
    X_cal : np.ndarray, shape (n_cal, n_features)
        Calibration feature matrix.
    y_cal : np.ndarray, shape (n_cal,)
        Calibration target values.

    Returns
    -------
    np.ndarray, shape (n_cal,)
        Array of non-negative nonconformity scores.
    """
    cal_preds  = model.predict(X_cal)    # Point predictions on calibration set
    scores     = np.abs(y_cal - cal_preds)  # Absolute residuals
    return scores


# Compute calibration nonconformity scores
cal_scores = compute_nonconformity_scores(
    model,
    splits_scaled['X_cal'],
    splits_scaled['y_cal']
)

print(f"Calibration nonconformity scores")
print(f"  Count  : {len(cal_scores)}")
print(f"  Mean   : {cal_scores.mean():.4f}")
print(f"  Median : {np.median(cal_scores):.4f}")
print(f"  90th % : {np.percentile(cal_scores, 90):.4f}")
print(f"  99th % : {np.percentile(cal_scores, 99):.4f}")

def calibrate_conformal_quantile(
    scores: np.ndarray,
    alpha: float
) -> float:
    """
    Compute the finite-sample-corrected conformal quantile q̂.

    The quantile level is adjusted by a factor of (1 + 1/n) to guarantee
    marginal coverage at level (1 - alpha) for finite calibration sets.

    Parameters
    ----------
    scores : np.ndarray
        Nonconformity scores from the calibration set.
    alpha : float
        Miscoverage level (e.g., 0.10 for 90% coverage intervals).

    Returns
    -------
    float
        Conformal quantile q̂ (interval half-width in target units).
    """
    n = len(scores)

    # Finite-sample-corrected quantile level
    # The (1 + 1/n) factor corrects for the fact that q̂ is estimated from data
    adjusted_level = min((1 - alpha) * (1 + 1 / n), 1.0)

    # np.quantile uses linear interpolation by default
    q_hat = np.quantile(scores, adjusted_level)
    return float(q_hat)


# Calibrate quantiles for multiple coverage levels
COVERAGE_LEVELS = [0.80, 0.90, 0.95, 0.99]
quantile_table  = {}

print(f"{'Coverage':>10s} | {'Alpha':>8s} | {'q̂ (half-width)':>18s}")
print("-" * 42)

for coverage in COVERAGE_LEVELS:
    alpha = 1.0 - coverage
    q_hat = calibrate_conformal_quantile(cal_scores, alpha)
    quantile_table[coverage] = q_hat
    print(f"  {coverage*100:5.0f}%   | {alpha:8.2f} | {q_hat:18.4f}")

def construct_prediction_intervals(
    model,
    X_test: np.ndarray,
    y_test: np.ndarray,
    q_hat: float
) -> pd.DataFrame:
    """
    Construct symmetric conformal prediction intervals for test observations.

    For each test point x_i, the interval is:
        [ŷ_i - q̂,  ŷ_i + q̂]

    Parameters
    ----------
    model : fitted regressor
        Trained base predictor.
    X_test : np.ndarray, shape (n_test, n_features)
        Test feature matrix.
    y_test : np.ndarray, shape (n_test,)
        True test targets.
    q_hat : float
        Conformal quantile (interval half-width).

    Returns
    -------
    pd.DataFrame
        DataFrame with columns: y_true, y_pred, lower, upper, covered, width.
    """
    y_pred = model.predict(X_test)   # Point predictions for all test observations

    lower = y_pred - q_hat           # Lower bound of conformal interval
    upper = y_pred + q_hat           # Upper bound of conformal interval

    # Indicator: True if the true value falls within the interval
    covered = (y_test >= lower) & (y_test <= upper)

    # Interval width (constant for split conformal with absolute-residual scores)
    width = upper - lower

    return pd.DataFrame({
        'y_true'  : y_test,
        'y_pred'  : y_pred,
        'lower'   : lower,
        'upper'   : upper,
        'covered' : covered,
        'width'   : width
    })


def evaluate_coverage(
    results_df: pd.DataFrame,
    nominal_coverage: float
) -> dict:
    """
    Compute empirical coverage and interval efficiency metrics.

    Parameters
    ----------
    results_df : pd.DataFrame
        Output of construct_prediction_intervals().
    nominal_coverage : float
        Target coverage level (e.g., 0.90).

    Returns
    -------
    dict
        Dictionary of coverage and efficiency metrics.
    """
    empirical_cov = results_df['covered'].mean()   # Fraction of covered test points
    mean_width    = results_df['width'].mean()      # Average interval width
    coverage_gap  = empirical_cov - nominal_coverage  # Should be >= 0 for valid method

    return {
        'nominal_coverage'  : nominal_coverage,
        'empirical_coverage': empirical_cov,
        'coverage_gap'      : coverage_gap,
        'mean_width'        : mean_width,
        'n_test'            : len(results_df)
    }


# Construct and evaluate intervals for all coverage levels
all_results = {}
all_metrics = {}

print(f"{'Coverage':>10s} | {'Empirical':>10s} | {'Gap':>8s} | {'Width':>10s} | {'Valid?':>7s}")
print("-" * 55)

for coverage in COVERAGE_LEVELS:
    q_hat   = quantile_table[coverage]
    results = construct_prediction_intervals(
        model, splits_scaled['X_test'], splits_scaled['y_test'], q_hat
    )
    metrics = evaluate_coverage(results, coverage)

    all_results[coverage] = results
    all_metrics[coverage] = metrics

    valid = '✓' if metrics['coverage_gap'] >= 0 else '✗'
    print(
        f"  {coverage*100:5.0f}%   |"
        f"  {metrics['empirical_coverage']*100:6.2f}%  |"
        f"  {metrics['coverage_gap']*100:+6.2f}% |"
        f"  {metrics['mean_width']:8.4f}  |"
        f"  {valid:>6s}"
    )

def summarise_coverage_tradeoff(all_metrics: dict) -> pd.DataFrame:
    """
    Assemble a summary DataFrame of coverage and efficiency metrics.

    Parameters
    ----------
    all_metrics : dict
        Dictionary mapping coverage level -> metrics dict.

    Returns
    -------
    pd.DataFrame
        Summary table with one row per coverage level.
    """
    rows = []
    for cov, m in all_metrics.items():
        rows.append({
            'Nominal Coverage (%)' : round(m['nominal_coverage'] * 100, 0),
            'Empirical Coverage (%)': round(m['empirical_coverage'] * 100, 2),
            'Coverage Gap (%)'     : round(m['coverage_gap'] * 100, 2),
            'Conformal q̂'         : round(quantile_table[cov], 4),
            'Mean Interval Width'  : round(m['mean_width'], 4),
        })
    return pd.DataFrame(rows)


# Generate and display the summary table
summary_df = summarise_coverage_tradeoff(all_metrics)
print("Coverage-Efficiency Summary")
print("=" * 70)
summary_df

def plot_conformal_diagnostics(
    cal_scores: np.ndarray,
    quantile_table: dict,
    all_results: dict,
    all_metrics: dict,
    primary_coverage: float = 0.90
) -> None:
    """
    Produce a 2x2 panel of conformal prediction diagnostic plots.

    Parameters
    ----------
    cal_scores : np.ndarray
        Calibration nonconformity scores.
    quantile_table : dict
        Mapping from coverage level to conformal quantile q̂.
    all_results : dict
        Mapping from coverage level to results DataFrame.
    all_metrics : dict
        Mapping from coverage level to metrics dictionary.
    primary_coverage : float
        Coverage level highlighted in the interval plot.
    """
    PALETTE = ['#2196F3', '#FF5722', '#4CAF50', '#9C27B0']  # Blue, Orange, Green, Purple

    fig = plt.figure(figsize=(16, 12))
    gs  = gridspec.GridSpec(2, 2, hspace=0.40, wspace=0.35)

    # ── Plot 1: Nonconformity Score Distribution ────────────────────────────────
    ax1 = fig.add_subplot(gs[0, 0])

    ax1.hist(cal_scores, bins=50, color='steelblue', alpha=0.75, edgecolor='white',
             label='Calibration scores')

    # Overlay vertical lines for each conformal quantile
    colors = PALETTE
    for i, (cov, q) in enumerate(quantile_table.items()):
        ax1.axvline(q, color=colors[i], linewidth=2, linestyle='--',
                    label=f'{int(cov*100)}% q̂ = {q:.3f}')

    ax1.set_xlabel('Nonconformity Score |y − ŷ|', fontsize=11)
    ax1.set_ylabel('Frequency', fontsize=11)
    ax1.set_title('Calibration Nonconformity Score Distribution', fontsize=12, fontweight='bold')
    ax1.legend(fontsize=9, loc='upper right')

    # ── Plot 2: Prediction Intervals on Test Set (primary coverage) ─────────────
    ax2 = fig.add_subplot(gs[0, 1])

    res  = all_results[primary_coverage]
    n_show = min(150, len(res))                          # Display first 150 test points
    idx  = np.arange(n_show)
    sub  = res.iloc[:n_show]

    # Shaded band: conformal interval
    ax2.fill_between(idx, sub['lower'], sub['upper'],
                     alpha=0.25, color='steelblue', label='90% CI band')

    # Point predictions
    ax2.plot(idx, sub['y_pred'], color='steelblue', linewidth=1.2, label='ŷ (prediction)', zorder=3)

    # True values: covered vs. uncovered
    covered_mask   = sub['covered'].values
    uncovered_mask = ~covered_mask

    ax2.scatter(idx[covered_mask],   sub['y_true'].values[covered_mask],
                color='green', s=12, alpha=0.6, label='y (covered)', zorder=4)
    ax2.scatter(idx[uncovered_mask], sub['y_true'].values[uncovered_mask],
                color='crimson', s=20, alpha=0.9, label='y (uncovered)', zorder=5)

    ax2.set_xlabel('Test Observation Index', fontsize=11)
    ax2.set_ylabel('Signal Value', fontsize=11)
    ax2.set_title(f'{int(primary_coverage*100)}% Conformal Prediction Intervals', fontsize=12, fontweight='bold')
    ax2.legend(fontsize=9)

    # ── Plot 3: Empirical vs. Nominal Coverage ───────────────────────────────────
    ax3 = fig.add_subplot(gs[1, 0])

    nominal   = [m['nominal_coverage'] * 100   for m in all_metrics.values()]
    empirical = [m['empirical_coverage'] * 100 for m in all_metrics.values()]
    x_pos     = np.arange(len(nominal))
    bar_width  = 0.35

    ax3.bar(x_pos - bar_width / 2, nominal,   bar_width, color='lightgray',    label='Nominal',   edgecolor='dimgray')
    ax3.bar(x_pos + bar_width / 2, empirical, bar_width, color=PALETTE[:len(nominal)], label='Empirical', edgecolor='white', alpha=0.85)

    ax3.set_xticks(x_pos)
    ax3.set_xticklabels([f'{n:.0f}%' for n in nominal])
    ax3.set_ylim(70, 103)
    ax3.set_xlabel('Nominal Coverage Level', fontsize=11)
    ax3.set_ylabel('Coverage (%)', fontsize=11)
    ax3.set_title('Empirical vs. Nominal Coverage', fontsize=12, fontweight='bold')
    ax3.legend(fontsize=10)

    # Annotate empirical values above bars
    for i, emp in enumerate(empirical):
        ax3.text(i + bar_width / 2, emp + 0.3, f'{emp:.1f}%', ha='center', va='bottom', fontsize=9)

    # ── Plot 4: Coverage vs. Interval Width Tradeoff ─────────────────────────────
    ax4 = fig.add_subplot(gs[1, 1])

    widths    = [m['mean_width'] for m in all_metrics.values()]

    color_cov   = 'steelblue'
    color_width = 'darkorange'

    lns1 = ax4.plot(nominal, empirical, 'o-', color=color_cov,   linewidth=2.0,
                    markersize=8, label='Empirical Coverage (%)')
    ax4.set_ylabel('Empirical Coverage (%)', color=color_cov, fontsize=11)
    ax4.tick_params(axis='y', labelcolor=color_cov)

    ax4b = ax4.twinx()    # Secondary y-axis for interval width
    lns2 = ax4b.plot(nominal, widths, 's--', color=color_width, linewidth=2.0,
                     markersize=8, label='Mean Interval Width')
    ax4b.set_ylabel('Mean Interval Width', color=color_width, fontsize=11)
    ax4b.tick_params(axis='y', labelcolor=color_width)

    ax4.set_xlabel('Nominal Coverage (%)', fontsize=11)
    ax4.set_title('Coverage–Efficiency Tradeoff', fontsize=12, fontweight='bold')

    # Combine legends from both axes
    all_lines  = lns1 + lns2
    all_labels = [l.get_label() for l in all_lines]
    ax4.legend(all_lines, all_labels, fontsize=9, loc='upper left')

    plt.suptitle(
        'Conformal Prediction Intervals for ML Signals — Diagnostic Report',
        fontsize=14, fontweight='bold', y=1.01
    )
    plt.savefig('conformal_prediction_diagnostics.png', bbox_inches='tight', dpi=150)
    plt.show()
    print("Diagnostic figure saved to: conformal_prediction_diagnostics.png")


# Render the diagnostic panel
plot_conformal_diagnostics(
    cal_scores,
    quantile_table,
    all_results,
    all_metrics,
    primary_coverage=0.90
)

def plot_adaptive_motivation(
    results_df: pd.DataFrame,
    coverage: float = 0.90
) -> None:
    """
    Plot interval width uniformity and residual heteroscedasticity.

    Visualises the constant-width limitation of split conformal prediction
    and highlights observations most likely to benefit from adaptive methods.

    Parameters
    ----------
    results_df : pd.DataFrame
        Output of construct_prediction_intervals() at the target coverage.
    coverage : float
        Nominal coverage level for annotation.
    """
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Compute absolute residuals for each test observation
    abs_residuals = np.abs(results_df['y_true'] - results_df['y_pred'])
    half_width    = (results_df['upper'] - results_df['lower']) / 2

    # ── Left panel: Residual vs. Prediction ──────────────────────────────────────
    ax = axes[0]
    scatter = ax.scatter(
        results_df['y_pred'],
        abs_residuals,
        c=results_df['covered'].astype(int),   # Colour by coverage status
        cmap='RdYlGn',
        alpha=0.5, s=15, edgecolors='none'
    )
    # Horizontal line at conformal quantile threshold
    ax.axhline(half_width.iloc[0], color='steelblue', linestyle='--', linewidth=1.8,
               label=f'q̂ = {half_width.iloc[0]:.3f}')
    ax.set_xlabel('Point Prediction ŷ', fontsize=11)
    ax.set_ylabel('Absolute Residual |y − ŷ|', fontsize=11)
    ax.set_title('Residual Magnitude vs. Prediction\n(Green = covered, Red = uncovered)', fontsize=11, fontweight='bold')
    ax.legend(fontsize=9)
    plt.colorbar(scatter, ax=ax, label='Covered (1) / Uncovered (0)')

    # ── Right panel: Empirical CDF of calibration scores ─────────────────────────
    ax2 = axes[1]
    sorted_scores = np.sort(cal_scores)
    cdf_values    = np.arange(1, len(sorted_scores) + 1) / len(sorted_scores)

    ax2.plot(sorted_scores, cdf_values, color='steelblue', linewidth=2, label='Empirical CDF')

    # Mark each quantile level on the CDF
    colors = ['#2196F3', '#FF5722', '#4CAF50', '#9C27B0']
    for i, (cov, q) in enumerate(quantile_table.items()):
        ax2.axvline(q, color=colors[i], linestyle=':', alpha=0.8, linewidth=1.5)
        ax2.axhline(cov, color=colors[i], linestyle=':', alpha=0.8, linewidth=1.5,
                    label=f'{int(cov*100)}% → q̂={q:.3f}')

    ax2.set_xlabel('Nonconformity Score', fontsize=11)
    ax2.set_ylabel('Empirical CDF', fontsize=11)
    ax2.set_title('Empirical CDF of Calibration Scores\nwith Conformal Quantiles', fontsize=11, fontweight='bold')
    ax2.legend(fontsize=8)

    plt.suptitle(
        'Constant-Width Limitation and Calibration CDF Analysis',
        fontsize=13, fontweight='bold'
    )
    plt.tight_layout()
    plt.savefig('conformal_adaptive_motivation.png', bbox_inches='tight', dpi=150)
    plt.show()
    print("Extended diagnostic figure saved to: conformal_adaptive_motivation.png")


# Render with the 90% coverage results
plot_adaptive_motivation(all_results[0.90], coverage=0.90)