Chapter 5 Experiment 5: Covariance Matrices and Block Matrix Operations

import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import warnings
warnings.filterwarnings('ignore')

np.set_printoptions(precision=4, suppress=True, linewidth=100)

# --- Color scheme (Accent Mix) ---
C_BG    = "#F8F8F8"
C_GRID  = "#D6D6D6"
C_AXIS  = "#000000"
C_V1    = "#57068C"   # violet
C_V2    = "#006385"   # deep blue
C_T1    = "#2AD2C9"   # teal
C_T2    = "#8900E1"   # Ultra Violet
C_WARN  = "#FF5D47"   # orange
C_AUX   = "#AB82C5"   # light violet
C_AREA  = "rgba(87, 6, 140, 0.25)"

def print_header(title):
    print("=" * 60)
    print(f"  {title}")
    print("=" * 60)

def print_step(step, desc):
    print(f"\n▶ Step {step}: {desc}")
    print("-" * 40)

def compare_print(label, actual, expected_desc):
    print(f"\n[{label}]")
    print(f"  Computed value:\n{actual}")
    print(f"  Expected value: {expected_desc}")

print_header("Experiment environment initialized")
print("  Color scheme: book palette")
print("  Plotting engine: Plotly")

# ============================================================
print_header("A1 | Loading the data")
# ============================================================

# --- Step 0 ---
X = np.array([
    [170, 65],
    [175, 70],
    [165, 60],
    [180, 75],
    [160, 55],
], dtype=float)

m, n = X.shape
labels = [f"Sample {k}" for k in range(m)]
feature_names = ["Height (cm)", "Weight (kg)"]

print_step(1, f"Data matrix X ∈ ℝ^{{{m}×{n}}}")
print(X)
print(f"\n  m={m} samples, n={n} features")

# ============================================================
print_header("A2 | Mean vector and centering")
# ============================================================

# --- Step 1 ---
x_bar = X.mean(axis=0)          # mean vector x̄ ∈ ℝ^n
print_step(1, "Compute the mean vector x̄ = (1/m) X^T 1_m")
print(f"  x̄ = {x_bar}")

# --- Step 2 ---
ones_m = np.ones((m, 1))
X_tilde = X - ones_m @ x_bar.reshape(1, -1)   # X̃ = X - 1_m x̄^T
print_step(2, "Centered data matrix X̃ = X - 1_m x̄^T")
print(X_tilde)

compare_print(
    "Column means (should all be 0)",
    X_tilde.mean(axis=0),
    "[0., 0.] (after centering, every feature has mean 0)"
)

# ============================================================
print_header("A3 | Interactive scatter plot")
# ============================================================

fig = make_subplots(
    rows=1, cols=2,
    subplot_titles=["Original data X", "Centered X̃"],
    horizontal_spacing=0.12
)

# --- Original data ---
for k in range(m):
    fig.add_trace(go.Scatter(
        x=[X[k, 0]], y=[X[k, 1]],
        mode='markers+text',
        marker=dict(size=12, color=C_V2),
        text=[labels[k]], textposition='top center',
        textfont=dict(size=11),
        name=labels[k], showlegend=False
    ), row=1, col=1)

# Mean point
fig.add_trace(go.Scatter(
    x=[x_bar[0]], y=[x_bar[1]],
    mode='markers+text',
    marker=dict(size=16, color=C_WARN, symbol='cross'),
    text=["x̄"], textposition='top right',
    textfont=dict(size=13, color=C_WARN),
    name="Mean x̄", showlegend=True
), row=1, col=1)

# --- After centering ---
for k in range(m):
    fig.add_trace(go.Scatter(
        x=[X_tilde[k, 0]], y=[X_tilde[k, 1]],
        mode='markers+text',
        marker=dict(size=12, color=C_V1),
        text=[labels[k]], textposition='top center',
        textfont=dict(size=11),
        name=labels[k], showlegend=False
    ), row=1, col=2)

# Origin (the mean moves to the origin)
fig.add_trace(go.Scatter(
    x=[0], y=[0],
    mode='markers+text',
    marker=dict(size=16, color=C_WARN, symbol='cross'),
    text=["Origin = x̄"], textposition='top right',
    textfont=dict(size=13, color=C_WARN),
    name="Origin", showlegend=True
), row=1, col=2)

# Coordinate-axis guide lines (dotted lines x=0 and y=0)
for col_idx, xs in enumerate([X, X_tilde], start=1):
    x_range = [xs[:, 0].min()-5, xs[:, 0].max()+5]
    fig.add_shape(type="line", x0=0, y0=xs[:, 1].min()-5,
                  x1=0, y1=xs[:, 1].max()+5,
                  line=dict(color=C_GRID, width=1, dash="dot"),
                  row=1, col=col_idx)
    fig.add_shape(type="line", x0=xs[:, 0].min()-5, y0=0,
                  x1=xs[:, 0].max()+5, y1=0,
                  line=dict(color=C_GRID, width=1, dash="dot"),
                  row=1, col=col_idx)

fig.update_layout(
    title=dict(text="Before and after centering: the mean vector moves to the origin", font=dict(size=15, color=C_AXIS)),
    plot_bgcolor=C_BG, paper_bgcolor=C_BG,
    height=450, width=850,
    font=dict(family="Microsoft JhengHei, SimHei, sans-serif")
)
fig.update_xaxes(gridcolor=C_GRID, gridwidth=0.8, title_text="Height (cm)")
fig.update_yaxes(gridcolor=C_GRID, gridwidth=0.8, title_text="Weight (kg)")
fig.show()
print("  Observation: in the right panel, the mean (×) has moved to the origin")

# ============================================================
print_header("B1 | Constructing the centering matrix H_m")
# ============================================================

I_m = np.eye(m)
ones_col = np.ones((m, 1))
H_m = I_m - (1/m) * (ones_col @ ones_col.T)

print_step(1, "H_m = I_m - (1/m) 1_m 1_m^T")
print(H_m)

# --- Verify: H_m 1_m = 0 ---
print_step(2, "Verify H_m 1_m = 0")
result = H_m @ ones_col.ravel()
compare_print("H_m 1_m", result, "[0, 0, ..., 0] (the zero vector)")

# --- Verify: H_m X = X̃ ---
print_step(3, "Verify H_m X = X̃")
X_tilde_check = H_m @ X
compare_print(
    "H_m X - X̃ (difference, should be all 0)",
    np.abs(X_tilde_check - X_tilde).max(),
    "0.0 (numerical error < 1e-12)"
)

# ============================================================
print_header("B2 | Experimenting with the centering matrix: student grade data")
# ============================================================

# --- Step 1: build the new data and H_4 ---
X_new = np.array([
    [80, 70, 90],
    [60, 80, 50],
    [90, 60, 70],
    [70, 90, 90],
], dtype=float)

m_new, n_new = X_new.shape
ones_new = np.ones((m_new, 1))
I_new    = np.eye(m_new)

# 🖊 Student fill-in: construct H_4
# Hint: follow B1, using I_new and ones_new
H_new = None  # change to: I_new - (1/m_new) * ones_new @ ones_new.T

print_step(1, "Apply H_4 to center X_new")
if H_new is not None:
    X_new_tilde = H_new @ X_new
    print("  X̃_new =")
    print(X_new_tilde)

    col_means = X_new_tilde.mean(axis=0)
    print(f"\n  Column means: {col_means}")
    all_zero = np.allclose(col_means, 0)
    print(f"  Are all the means 0? {'✓ yes' if all_zero else '✗ no; check how H_new is written'}")
else:
    print("  [Fill in H_new first]")

# --- Step 2: the constant-column experiment ---
print_step(2, "Append a constant column (all 100) and observe the effect of centering")

# 🖊 Student fill-in: append a column of all 100s on the right of X_new
# Hint: np.hstack([X_new, np.full((m_new, 1), 100.0)])
X_aug = None  # change to the hint above

if X_aug is not None and H_new is not None:
    X_aug_tilde = H_new @ X_aug
    print("  X_aug (with the constant column) =")
    print(X_aug)
    print("\n  H_4 @ X_aug (after centering) =")
    print(X_aug_tilde)
    const_col = X_aug_tilde[:, -1]
    print(f"\n  Constant column after centering: {const_col}")
    print(f"  → after centering, the constant column is {'all 0 ✓' if np.allclose(const_col, 0) else 'not 0 ✗'}")
    print("  Interpretation: a constant column carries no information about variation between samples;")
    print("        it vanishes after centering → it contributes nothing to the covariance matrix.")
else:
    print("  [Fill in X_aug and H_new first]")

# --- Step 3: compute and interpret the covariance matrix ---
print_step(3, "Compute S_new and interpret its entries")

# 🖊 Student fill-in: compute S_new from the definition form
# Hint: first make sure X_new_tilde has been computed, then use (1/(m_new-1)) * X_new_tilde.T @ X_new_tilde
S_new = None  # change to the hint above

subject_names = ["Math", "English", "PE"]
if S_new is not None:
    print("\n  S_new =")
    print(S_new)
    print("\n  Diagonal entries (variance of each subject):")
    for j, name in enumerate(subject_names):
        print(f"    Var({name}) = S[{j},{j}] = {S_new[j,j]:.2f}")

    print("\n  Off-diagonal entries (covariances between subjects):")
    for i in range(n_new):
        for j in range(i+1, n_new):
            cov_ij = S_new[i, j]
            sign = "positive correlation ↑↑" if cov_ij > 50 else ("negative correlation ↑↓" if cov_ij < -50 else "nearly unrelated")
            print(f"    Cov({subject_names[i]}, {subject_names[j]}) = {cov_ij:.2f}  → {sign}")

    # Find the strongest correlation
    S_abs = np.abs(S_new - np.diag(np.diag(S_new)))
    i_max, j_max = np.unravel_index(np.argmax(S_abs), S_new.shape)
    print(f"\n  The two most strongly correlated subjects: {subject_names[i_max]} and {subject_names[j_max]}")
    print(f"  Cov = {S_new[i_max, j_max]:.2f}")
else:
    print("  [Fill in S_new first]")

# ============================================================
print_header("B3 | Heat map of the matrix H_m")
# ============================================================

fig = make_subplots(
    rows=1, cols=3,
    subplot_titles=[
        "I_m (identity matrix)",
        f"(1/m) 1_m 1_m^T (mean projection)",
        "H_m = I_m - (1/m) 1_m 1_m^T"
    ],
    horizontal_spacing=0.08
)

rank1 = (1/m) * (ones_col @ ones_col.T)
matrices = [I_m, rank1, H_m]
colormaps = ['Blues', 'Oranges', 'Purples']

for i, (mat, cmap) in enumerate(zip(matrices, colormaps)):
    colorscale_map = {
        'Blues':   [[0, C_BG], [1, C_V2]],
        'Oranges': [[0, C_BG], [1, C_WARN]],
        'Purples': [
            [0.0,  C_WARN],
            [0.4,  C_BG],
            [1.0,  C_V1]
        ]
    }
    fig.add_trace(go.Heatmap(
        z=mat,
        colorscale=colorscale_map[cmap],
        zmin=-0.3, zmax=1.0,
        text=[[f"{v:.2f}" for v in row] for row in mat],
        texttemplate="%{text}",
        textfont=dict(size=11),
        showscale=(i == 2),
        colorbar=dict(x=1.02)
    ), row=1, col=i+1)

fig.update_layout(
    title=dict(text=f"Structure of H_{m}: the identity matrix minus the mean projection", font=dict(size=14, color=C_AXIS)),
    plot_bgcolor=C_BG, paper_bgcolor=C_BG,
    height=380, width=860,
    font=dict(family="Microsoft JhengHei, SimHei, sans-serif")
)
fig.update_yaxes(autorange="reversed")
fig.show()
print(f"  Diagonal entries = {1 - 1/m:.4f} = 1 - 1/m")
print(f"  Off-diagonal entries = {-1/m:.4f} = -1/m")

# ============================================================
print_header("C1 | The covariance matrix: checking the three forms side by side + active exploration")
# ============================================================

# ── Reuse the height/weight data of Block A ──────────────────────────────
print_step(1, "Definition form: S = (1/(m-1)) X̃^T X̃")
S_def = (1/(m-1)) * X_tilde.T @ X_tilde
print(S_def)

print_step(2, "Operator form: S = (1/(m-1)) X^T H_m X")
# 🖊 Student fill-in ↓
S_op = None   # change to: (1/(m-1)) * X.T @ H_m @ X

print_step(3, "Expanded form: S = (1/(m-1)) (X^T X - m x̄ x̄^T)")
# 🖊 Student fill-in ↓
S_exp = None  # change to: (1/(m-1)) * (X.T @ X - m * np.outer(x_bar, x_bar))

# --- Check the three forms ---
print("\n" + "="*40)
print("  Check: largest absolute difference among the three forms")
print("="*40)
if S_op is not None:
    compare_print("definition form vs operator form", np.abs(S_def - S_op).max(), "0.0")
else:
    print("  [Fill in step 2 first]")
if S_exp is not None:
    compare_print("definition form vs expanded form", np.abs(S_def - S_exp).max(), "0.0")
else:
    print("  [Fill in step 3 first]")

print(f"\n  S = \n{S_def}")
print(f"  S_{{00}} = {S_def[0,0]:.2f}  → variance of height")
print(f"  S_{{11}} = {S_def[1,1]:.2f}  → variance of weight")
print(f"  S_{{01}} = {S_def[0,1]:.2f}  → positive correlation (the taller, the heavier)")

# ── Active exploration: how does the covariance matrix change when a constant feature is added? ───────────
print_step(4, "Active exploration: how S changes after adding a constant feature (all 170)")
print("  Construct X_const = [X | 170 * 1_m]")

X_const = np.hstack([X, np.full((m, 1), 170.0)])

# 🖊 Student prediction: what will row 2 and column 2 of S_const be?
# Write your prediction below first, then run the computation that follows to check it
# Prediction: column/row 2 of S_const (0-based) = ____________

x_bar_const = X_const.mean(axis=0)
ones_col_c  = np.ones((m, 1))
H_const     = np.eye(m) - (1/m) * ones_col_c @ ones_col_c.T
X_const_tilde = H_const @ X_const
S_const = (1/(m-1)) * X_const_tilde.T @ X_const_tilde

print("\n  S_const (3×3) =")
print(S_const)
print("\n  Column/row 2 (the constant feature):")
print(f"    {S_const[2, :]}")
print("  → The constant feature has variance 0, and its covariance with every feature is also 0.")
print("  → Geometrically: after centering, the constant feature is the zero vector and contributes nothing to X̃^T X̃.")

# Colab: enable the custom widget manager so that FigureWidget and the sliders are interactive; skipped elsewhere
try:
    from google.colab import output
    output.enable_custom_widget_manager()
except ImportError:
    pass

# ============================================================
print_header("C2 | Exploring correlation: fixed typical scenarios + interactive generation")
# ============================================================

import ipywidgets as widgets
from IPython.display import display

# ------------------------------------------------------------------
# Utility: generate bivariate data with a given correlation coefficient exactly, via Cholesky
# ------------------------------------------------------------------
def generate_correlated(r, m_samples, sigma=1.0, seed=0):
    # Generate bivariate normal data with correlation r via Cholesky decomposition.
    # S = sigma^2 * [[1, r],[r, 1]] = L L^T
    # L = sigma * [[1, 0],[r, sqrt(1-r^2)]]
    rng = np.random.default_rng(seed)
    Z = rng.standard_normal((m_samples, 2))      # Z ~ N(0, I)
    L = sigma * np.array([[1.0, 0.0],
                           [r,  np.sqrt(1 - r**2)]])
    X_raw = Z @ L.T                               # X = Z L^T
    # center by hand so that the sample mean is exactly 0
    X_centered = X_raw - X_raw.mean(axis=0)
    return X_centered

def compute_cov_and_r(data):
    # Compute sample covariance matrix and Pearson r.
    S = np.cov(data.T, ddof=1)
    r_actual = S[0, 1] / np.sqrt(S[0, 0] * S[1, 1])
    return S, r_actual

# ------------------------------------------------------------------
# Part 1: static comparison of the three typical scenarios
# ------------------------------------------------------------------
print_step(1, "Fixed typical scenarios: strong positive / no / strong negative correlation")

SCENARIOS = [
    dict(r_target=0.90,  label="Strong positive correlation",  color=C_V1,   m_s=80, seed=7),
    dict(r_target=0.00,  label="No correlation",    color=C_T1,   m_s=80, seed=7),
    dict(r_target=-0.90, label="Strong negative correlation",  color=C_WARN, m_s=80, seed=7),
]

fig_static = make_subplots(
    rows=1, cols=3,
    horizontal_spacing=0.10,
    subplot_titles=[sc["label"] for sc in SCENARIOS]
)

for col_i, sc in enumerate(SCENARIOS, start=1):
    D = generate_correlated(sc["r_target"], sc["m_s"], sigma=1.0, seed=sc["seed"])
    S_sc, r_act = compute_cov_and_r(D)

    # Scatter points
    fig_static.add_trace(go.Scatter(
        x=D[:, 0], y=D[:, 1],
        mode='markers',
        marker=dict(size=5, color=sc["color"], opacity=0.65),
        showlegend=False,
        hovertemplate=(
            "x₀=%{x:.3f}<br>x₁=%{y:.3f}"
            f"<extra>{sc['label']}</extra>"
        )
    ), row=1, col=col_i)

    # Mean point
    fig_static.add_trace(go.Scatter(
        x=[D[:, 0].mean()], y=[D[:, 1].mean()],
        mode='markers',
        marker=dict(size=10, color=C_AXIS, symbol='cross'),
        showlegend=False,
        hovertemplate="Mean x̄<extra></extra>"
    ), row=1, col=col_i)

    # Annotate the matrix entries
    annotation_text = (
        f"r_target = {sc['r_target']:+.2f}<br>"
        f"r_actual = {r_act:+.4f}<br>"
        f"S₀₀ = {S_sc[0,0]:.4f}<br>"
        f"S₁₁ = {S_sc[1,1]:.4f}<br>"
        f"S₀₁ = {S_sc[0,1]:+.4f}"
    )
    x_anchor = D[:, 0].min() + 0.05
    y_anchor = D[:, 1].max() - 0.1
    fig_static.add_annotation(
        text=annotation_text,
        xref=f"x{col_i}", yref=f"y{col_i}",
        x=x_anchor, y=y_anchor,
        showarrow=False,
        font=dict(size=10, color=sc["color"]),
        align="left", bgcolor=C_BG,
        bordercolor=sc["color"], borderwidth=1
    )

    print(f"  {sc['label']:5s}: r_target={sc['r_target']:+.2f}  "
          f"r_actual={r_act:+.4f}  "
          f"S = [[{S_sc[0,0]:.3f}, {S_sc[0,1]:+.3f}],"
          f"[{S_sc[1,0]:+.3f}, {S_sc[1,1]:.3f}]]")

fig_static.update_layout(
    title=dict(
        text="Exact generation via Cholesky: three typical correlation structures (m=80, σ=1)",
        font=dict(size=14, color=C_AXIS)
    ),
    plot_bgcolor=C_BG, paper_bgcolor=C_BG,
    height=430, width=920,
    font=dict(family="Microsoft JhengHei, SimHei, sans-serif")
)
fig_static.update_xaxes(
    gridcolor=C_GRID, gridwidth=0.8,
    zeroline=True, zerolinecolor=C_AXIS, zerolinewidth=1,
    title_text="Feature x₀"
)
fig_static.update_yaxes(
    gridcolor=C_GRID, gridwidth=0.8,
    zeroline=True, zerolinecolor=C_AXIS, zerolinewidth=1,
    scaleanchor="x", scaleratio=1,
    title_text="Feature x₁"
)
fig_static.show()

# ------------------------------------------------------------------
# Part 2: interactive exploration (ipywidgets + FigureWidget)
# ------------------------------------------------------------------
print_step(2, "Interactive exploration: adjust r, m, σ and watch the covariance matrix change")
print("  ↓ The figure is redrawn automatically when you move a slider")

# ============================================================
# Interactive widget — adjust r / m / σ and update the figure and the covariance matrix in real time
# ============================================================

def make_interactive_scatter():
    # Build interactive scatter plot + covariance matrix display.

    # --- Initial parameters ---
    r_init, m_init, sigma_init = 0.70, 60, 1.0

    # --- FigureWidget (updated in place, without re-rendering the whole page) ---
    D_init = generate_correlated(r_init, m_init, sigma=sigma_init, seed=42)
    S_init, r_init_act = compute_cov_and_r(D_init)

    scatter_trace = go.Scatter(
        x=D_init[:, 0], y=D_init[:, 1],
        mode='markers',
        marker=dict(size=5, color=C_V1, opacity=0.65),
        name="Data points"
    )
    mean_trace = go.Scatter(
        x=[0], y=[0],
        mode='markers',
        marker=dict(size=12, color=C_AXIS, symbol='cross'),
        name="Mean"
    )
    fig_w = go.FigureWidget(
        data=[scatter_trace, mean_trace],
        layout=go.Layout(
            title=dict(text="Interactive: adjust the parameters and watch the covariance structure", font=dict(size=13, color=C_AXIS)),
            plot_bgcolor=C_BG, paper_bgcolor=C_BG,
            xaxis=dict(
                title="Feature x₀", gridcolor=C_GRID, gridwidth=0.8,
                zeroline=True, zerolinecolor=C_AXIS,
                range=[-4, 4]
            ),
            yaxis=dict(
                title="Feature x₁", gridcolor=C_GRID, gridwidth=0.8,
                zeroline=True, zerolinecolor=C_AXIS,
                scaleanchor='x', scaleratio=1,
                range=[-4, 4]
            ),
            height=400, width=550,
            font=dict(family="Microsoft JhengHei, SimHei, sans-serif"),
            showlegend=True
        )
    )

    # --- Output text box ---
    out_text = widgets.HTML(value=_format_cov_html(S_init, r_init_act, r_init, m_init, sigma_init))

    # --- Sliders ---
    slider_r = widgets.FloatSlider(
        value=r_init, min=-0.99, max=0.99, step=0.01,
        description='Correlation r:',
        style={'description_width': '110px'},
        layout=widgets.Layout(width='500px'),
        continuous_update=True
    )
    slider_m = widgets.IntSlider(
        value=m_init, min=10, max=300, step=10,
        description='Samples m:',
        style={'description_width': '110px'},
        layout=widgets.Layout(width='500px'),
        continuous_update=False
    )
    slider_sigma = widgets.FloatSlider(
        value=sigma_init, min=0.5, max=3.0, step=0.1,
        description='Std. dev. σ:',
        style={'description_width': '110px'},
        layout=widgets.Layout(width='500px'),
        continuous_update=True
    )
    seed_box = widgets.BoundedIntText(
        value=42, min=0, max=9999,
        description='Random seed:',
        style={'description_width': '110px'},
        layout=widgets.Layout(width='200px')
    )

    def _update(change):
        r_val    = slider_r.value
        m_val    = slider_m.value
        sigma_val = slider_sigma.value
        seed_val  = seed_box.value

        D_new = generate_correlated(r_val, m_val, sigma=sigma_val, seed=seed_val)
        S_new, r_new = compute_cov_and_r(D_new)

        lim = max(3.5, float(np.abs(D_new).max()) * 1.15)

        with fig_w.batch_update():
            fig_w.data[0].x = D_new[:, 0]
            fig_w.data[0].y = D_new[:, 1]
            fig_w.data[1].x = [D_new[:, 0].mean()]
            fig_w.data[1].y = [D_new[:, 1].mean()]
            fig_w.layout.xaxis.range = [-lim, lim]
            fig_w.layout.yaxis.range = [-lim, lim]

        out_text.value = _format_cov_html(S_new, r_new, r_val, m_val, sigma_val)

    slider_r.observe(_update, names='value')
    slider_m.observe(_update, names='value')
    slider_sigma.observe(_update, names='value')
    seed_box.observe(_update, names='value')

    # --- Layout ---
    controls = widgets.VBox([slider_r, slider_m, slider_sigma, seed_box])
    right_panel = widgets.VBox([out_text])
    display(widgets.HBox([
        widgets.VBox([fig_w, controls]),
        widgets.VBox([right_panel])
    ]))


def _format_cov_html(S, r_act, r_target, m_val, sigma_val):
    # Format covariance matrix as HTML for widget display.
    color_r = C_V1 if r_act >= 0 else C_WARN
    eigenvalues = np.linalg.eigvalsh(S)[::-1]

    lines = [
        f'<div style="font-family:monospace;font-size:13px;padding:12px;'
        f'background:{C_BG};border:1.5px solid {C_V1};border-radius:6px;'
        f'min-width:280px;margin-left:20px;margin-top:10px;">',
        f'<b style="color:{C_V1};font-size:14px;">Covariance matrix S</b><br><br>',
        f'<table style="border-collapse:collapse;margin-bottom:8px;">',
        f'<tr>',
        f'  <td style="padding:4px 10px;border:1px solid {C_GRID};">S₀₀ = {S[0,0]:.4f}</td>',
        f'  <td style="padding:4px 10px;border:1px solid {C_GRID};color:{C_V2};">S₀₁ = {S[0,1]:+.4f}</td>',
        f'</tr><tr>',
        f'  <td style="padding:4px 10px;border:1px solid {C_GRID};color:{C_V2};">S₁₀ = {S[1,0]:+.4f}</td>',
        f'  <td style="padding:4px 10px;border:1px solid {C_GRID};">S₁₁ = {S[1,1]:.4f}</td>',
        f'</tr></table>',
        f'<hr style="border:none;border-top:1px solid {C_GRID};margin:6px 0;">',
        f'<b>Pearson correlation coefficient</b><br>',
        f'&nbsp; r target = {r_target:+.2f}<br>',
        f'&nbsp; r actual = <span style="color:{color_r};font-weight:bold;">{r_act:+.4f}</span><br><br>',
        f'<span style="color:#888;font-size:11px;">◆ Eigenvalues (variances of the principal components)—understandable after §8</span><br>',
        f'<span style="color:#aaa;font-size:11px;">&nbsp; λ₀ ≈ {eigenvalues[0]:.4f}, λ₁ ≈ {eigenvalues[1]:.4f}</span><br><br>',
        f'<b>Trace tr(S)</b> = {np.trace(S):.4f}<br>',
        f'<b>σ setting</b> = {sigma_val:.1f} &nbsp;→&nbsp; <b>tr(S)/n</b> ≈ {np.trace(S)/2:.4f}<br>',
        f'<b>Number of samples m</b> = {m_val}',
        f'</div>',
    ]
    return '\n'.join(lines)

make_interactive_scatter()

# ============================================================
print_header("D1 | Total variance: by feature vs. by sample")
# ============================================================

print_step(1, "Summing over features: tr(S)")
total_var_feature = np.trace(S_def)
print(f"  tr(S) = {S_def[0,0]:.4f} + {S_def[1,1]:.4f} = {total_var_feature:.4f}")

print_step(2, "Summing over samples: (1/(m-1)) Σ ||x̃^(k)||²")
norms_sq = np.sum(X_tilde**2, axis=1)    # ||x̃^(k)||² for each k
total_var_sample = norms_sq.sum() / (m - 1)
print(f"  ||x̃^(k)||² for each sample: {norms_sq}")
print(f"  sum / (m-1) = {norms_sq.sum():.4f} / {m-1} = {total_var_sample:.4f}")

print_step(3, "Frobenius norm: (1/(m-1)) ||X̃||_F²")
frob_sq = np.sum(X_tilde**2)
total_var_frob = frob_sq / (m - 1)
print(f"  ||X̃||_F² = {frob_sq:.4f}")
print(f"  / (m-1) = {total_var_frob:.4f}")

print_step(4, "Cyclic property of the trace: tr(X̃^T X̃) vs tr(X̃ X̃^T)")
trace_XtX  = np.trace(X_tilde.T @ X_tilde)
trace_XXt  = np.trace(X_tilde @ X_tilde.T)
compare_print(
    "tr(X̃^T X̃) == tr(X̃ X̃^T)",
    f"{trace_XtX:.4f} vs {trace_XXt:.4f}",
    "equal (cyclic property of the trace)"
)

print("\n" + "="*40)
print(f"  All three computations give: {total_var_feature:.4f}")
print("="*40)

# ============================================================
print_header("D2 (exercise) | Estimate the spread radius from tr(S), then check with the data")
# ============================================================

# --- Step 1: estimate the standard deviation of each feature ---
trace_S_given = 48.0
n_features    = 3

print_step(1, "Estimate the standard deviation of each feature")
# 🖊 Student fill-in: tr(S) = Var(x_0) + Var(x_1) + Var(x_2)
#    under the assumption of similar variances, the variance of each feature ≈ tr(S) / n
#    standard deviation = sqrt(variance of each feature)
sigma_hat = None  # change to: np.sqrt(trace_S_given / n_features)

if sigma_hat is not None:
    print(f"  Estimated variance of each feature ≈ {trace_S_given}/{n_features} = {trace_S_given/n_features:.2f}")
    print(f"  Estimated standard deviation σ̂ ≈ {sigma_hat:.4f}")
else:
    print("  [Fill in sigma_hat]")

# --- Step 2: estimate the squared Frobenius norm ---
print_step(2, "Estimate ||X̃||_F²")
m_given = 20
# 🖊 Student fill-in: ||X̃||_F² = (m-1) * tr(S)
frob_sq_hat = None  # change to: (m_given - 1) * trace_S_given

if frob_sq_hat is not None:
    print(f"  ||X̃||_F² ≈ (m-1) × tr(S) = {m_given-1} × {trace_S_given} = {frob_sq_hat:.1f}")
else:
    print("  [Fill in frob_sq_hat]")

# --- Step 3: check with the full data ---
print_step(3, "Check with the full data")

# Generate data with a fixed seed: the population variances sum to 3.8²+4.2²+4.0² = 48.08 ≈ 48, matching the report;
# but with only m=20 samples, the sample tr(S) has sampling error (for this seed, tr(S) ≈ 37.55 and the feature σ ≈ 3.19, 4.19, 3.13)
rng = np.random.default_rng(99)
# the population standard deviations of the three features are 3.8, 4.2, 4.0 (average about 4, close to sqrt(16)≈4)
L_true = np.diag([3.8, 4.2, 4.0])
Z_true = rng.standard_normal((m_given, n_features))
X_verify = Z_true @ L_true + np.array([50.0, 60.0, 70.0])  # add a nonzero mean

# Compute the actual values
x_bar_v   = X_verify.mean(axis=0)
ones_v    = np.ones((m_given, 1))
H_v       = np.eye(m_given) - (1/m_given) * ones_v @ ones_v.T
X_v_tilde = H_v @ X_verify
S_verify  = (1/(m_given-1)) * X_v_tilde.T @ X_v_tilde

trace_actual  = np.trace(S_verify)
frob_sq_actual = np.sum(X_v_tilde**2)
sigma_actual  = np.sqrt(np.diag(S_verify))

print(f"  Actual tr(S)     = {trace_actual:.4f}  (reported: {trace_S_given})")
print(f"  Actual ||X̃||_F² = {frob_sq_actual:.4f}  (estimate: {frob_sq_hat})")
print(f"  Actual σ of each feature = {sigma_actual}")
if sigma_hat is not None:
    print(f"  Estimated σ̂     = {sigma_hat:.4f} (common estimate)")
print(f"\n  Check: (m-1)*tr(S) = {(m_given-1)*trace_actual:.4f} ≈ ||X̃||_F² = {frob_sq_actual:.4f}")
print(f"  Difference: {abs((m_given-1)*trace_actual - frob_sq_actual):.2e} (numerical error)")

# ============================================================
print_header("D3 | Geometric picture of the Frobenius norm")
# ============================================================

fig = go.Figure()

# Background grid lines and origin
fig.add_hline(y=0, line_color=C_GRID, line_width=1)
fig.add_vline(x=0, line_color=C_GRID, line_width=1)

colors_pts = [C_V1, C_V2, C_T1, C_T2, C_WARN]

for k in range(m):
    xk = X_tilde[k, 0]
    yk = X_tilde[k, 1]
    norm_sq = xk**2 + yk**2
    norm = np.sqrt(norm_sq)
    col = colors_pts[k]

    # Vector (from the origin to the sample point)
    fig.add_annotation(
        x=xk, y=yk, ax=0, ay=0,
        xref='x', yref='y', axref='x', ayref='y',
        showarrow=True,
        arrowhead=3, arrowsize=1.2, arrowwidth=2,
        arrowcolor=col
    )

    # Sample point
    fig.add_trace(go.Scatter(
        x=[xk], y=[yk],
        mode='markers+text',
        marker=dict(size=12, color=col),
        text=[f"x̃⁽{k}⁾<br>‖‖²={norm_sq:.1f}"],
        textposition='top right' if xk >= 0 else 'top left',
        textfont=dict(size=10, color=col),
        name=f"Sample {k} (‖x̃‖={norm:.2f})",
        hovertemplate=(
            f"x̃⁽{k}⁾ = ({xk:.1f}, {yk:.1f})<br>"
            f"‖x̃⁽{k}⁾‖² = {norm_sq:.2f}"
            "<extra></extra>"
        )
    ))

# Origin
fig.add_trace(go.Scatter(
    x=[0], y=[0], mode='markers',
    marker=dict(size=14, color=C_AXIS, symbol='cross'),
    name="Origin (mean)", showlegend=True
))

# Annotate the Frobenius norm
frob_sq_val = np.sum(X_tilde**2)
fig.add_annotation(
    x=0.02, y=0.98, xref='paper', yref='paper',
    text=(
        f"‖X̃‖²_F = Σ ‖x̃⁽ᵏ⁾‖² = <b>{frob_sq_val:.2f}</b><br>"
        f"Total variance = ‖X̃‖²_F / (m-1) = {frob_sq_val/(m-1):.2f}"
    ),
    showarrow=False,
    font=dict(size=12, color=C_V1),
    align="left", bgcolor=C_BG,
    bordercolor=C_V1, borderwidth=1.5
)

fig.update_layout(
    title=dict(
        text="Centered sample vectors: ‖X̃‖²_F = sum of the squared lengths of all the vectors",
        font=dict(size=14, color=C_AXIS)
    ),
    plot_bgcolor=C_BG, paper_bgcolor=C_BG,
    xaxis=dict(title="Feature x̃₀ (height deviation)", gridcolor=C_GRID,
               gridwidth=0.8, zeroline=True, zerolinecolor=C_AXIS),
    yaxis=dict(title="Feature x̃₁ (weight deviation)", gridcolor=C_GRID,
               gridwidth=0.8, zeroline=True, zerolinecolor=C_AXIS,
               scaleanchor='x', scaleratio=1),
    height=500, width=620,
    font=dict(family="Microsoft JhengHei, SimHei, sans-serif"),
    showlegend=True,
    legend=dict(x=1.02, y=1, bgcolor=C_BG, bordercolor=C_GRID)
)
fig.show()

# ============================================================
print_header("D4 | Heat map of the covariance matrix + preview of eigenvalues")
# ============================================================

eigenvalues, eigenvectors = np.linalg.eigh(S_def)
idx = np.argsort(eigenvalues)[::-1]   # sort from largest to smallest
eigenvalues = eigenvalues[idx]
eigenvectors = eigenvectors[:, idx]

print_step(1, "Eigenvalues (= variances of the principal components)")
for i, lam in enumerate(eigenvalues):
    pct = 100 * lam / eigenvalues.sum()
    print(f"  λ_{i} = {lam:.4f}  (explains {pct:.1f}% of the total variance)")

print_step(2, "Check: tr(S) = Σ λ_i")
compare_print("Σ λ_i", eigenvalues.sum(), f"tr(S) = {np.trace(S_def):.4f}")

# --- Figure 1: heat map of S ---
fig = make_subplots(
    rows=1, cols=2,
    subplot_titles=["Covariance matrix S (heat map)", "Principal directions (eigenvectors)"],
    horizontal_spacing=0.15
)

fig.add_trace(go.Heatmap(
    z=S_def,
    colorscale=[[0, C_BG], [1, C_V1]],
    text=[[f"{v:.2f}" for v in row] for row in S_def],
    texttemplate="%{text}",
    textfont=dict(size=16),
    showscale=True,
    colorbar=dict(x=0.45),
    xgap=2, ygap=2
), row=1, col=1)

fig.update_xaxes(
    tickvals=[0,1], ticktext=["x₀(height)", "x₁(weight)"],
    row=1, col=1
)
fig.update_yaxes(
    tickvals=[0,1], ticktext=["x₀(height)", "x₁(weight)"],
    autorange='reversed', row=1, col=1
)

# --- Figure 2: scatter points + eigenvector directions ---
scale_factors = [np.sqrt(lam) * 2.5 for lam in eigenvalues]
evec_colors = [C_V1, C_T1]

for k in range(m):
    fig.add_trace(go.Scatter(
        x=[X_tilde[k,0]], y=[X_tilde[k,1]],
        mode='markers',
        marker=dict(size=9, color=C_V2, opacity=0.8),
        showlegend=False
    ), row=1, col=2)

for i in range(n):
    ev = eigenvectors[:, i]
    s  = scale_factors[i]
    fig.add_annotation(
        x=ev[0]*s, y=ev[1]*s, ax=0, ay=0,
        xref='x2', yref='y2', axref='x2', ayref='y2',
        showarrow=True,
        arrowhead=3, arrowsize=1.3, arrowwidth=3,
        arrowcolor=evec_colors[i]
    )
    fig.add_annotation(
        x=ev[0]*s*1.15, y=ev[1]*s*1.15,
        xref='x2', yref='y2',
        text=f"<b>PC{i}</b><br>λ={eigenvalues[i]:.1f}",
        showarrow=False,
        font=dict(size=11, color=evec_colors[i])
    )

fig.update_xaxes(
    gridcolor=C_GRID, gridwidth=0.8, zeroline=True, zerolinecolor=C_AXIS,
    title_text="x̃₀", row=1, col=2
)
fig.update_yaxes(
    gridcolor=C_GRID, gridwidth=0.8, zeroline=True, zerolinecolor=C_AXIS,
    scaleanchor='x2', scaleratio=1, title_text="x̃₁", row=1, col=2
)

fig.update_layout(
    title=dict(
        text="Covariance matrix and principal directions (Chapter 11: preview of PCA)",
        font=dict(size=14, color=C_AXIS)
    ),
    plot_bgcolor=C_BG, paper_bgcolor=C_BG,
    height=460, width=880,
    font=dict(family="Microsoft JhengHei, SimHei, sans-serif")
)
fig.show()
print("  → PC0 direction = direction of maximal variance of the data = first principal axis; PC1 is orthogonal to it, the direction of minimal variance")
print("  → See Chapter 11: Singular Value Decomposition and PCA")

# ============================================================
print_header("E | Capstone task: complete analysis of three-dimensional physiological data")
# ============================================================

# Data (do not peek at the answer for S!)
X_task = np.array([
    [210, 42, 72],
    [185, 55, 65],
    [230, 38, 80],
    [195, 60, 68],
    [220, 45, 75],
    [175, 58, 62],
    [240, 35, 82],
    [200, 50, 70],
], dtype=float)

feature_names_task = ["Reaction time(ms)", "Grip strength(kg)", "Heart rate(bpm)"]
m_task, n_task = X_task.shape
print(f"  Data size: m={m_task} samples, n={n_task} features")
print("  X_task =")
print(X_task)

# ══════════════════════════════════════════════════════════
# 🖊 Step 1: compute the mean vector x̄
# ══════════════════════════════════════════════════════════
print_step(1, "Compute the mean vector x̄")
# Fill in the code ↓ (hint: X_task.mean(axis=?))
x_bar_task = None  # change to: X_task.mean(axis=0)

if x_bar_task is not None:
    print(f"  x̄ = {x_bar_task}")
    for j, name in enumerate(feature_names_task):
        print(f"    mean of {name} = {x_bar_task[j]:.2f}")
else:
    print("  [Fill in x_bar_task]")

# ══════════════════════════════════════════════════════════
# 🖊 Step 2: construct H_8 and compute the centered matrix X̃
# ══════════════════════════════════════════════════════════
print_step(2, "Construct H_8 and compute X̃_task = H_8 @ X_task")
# Fill in the code ↓
ones_task   = np.ones((m_task, 1))
H_task      = None  # change to: np.eye(m_task) - (1/m_task) * ones_task @ ones_task.T
X_task_tilde = None  # change to: H_task @ X_task

if X_task_tilde is not None:
    print("  X̃_task =")
    print(X_task_tilde)
    col_means_task = X_task_tilde.mean(axis=0)
    ok = np.allclose(col_means_task, 0)
    print(f"  Column means (should be 0): {col_means_task.round(10)}")
    print(f"  Check: {'✓ passed' if ok else '✗ check H_task'}")
else:
    print("  [Fill in H_task and X_task_tilde]")

# ══════════════════════════════════════════════════════════
# 🖊 Step 3: compute the covariance matrix S_task (with the expanded form)
# ══════════════════════════════════════════════════════════
print_step(3, "Compute the covariance matrix S_task (expanded form)")
# Expanded form: S = (1/(m-1)) * (X^T X - m * x̄ x̄^T)
# Fill in the code ↓
S_task = None  # change to: (1/(m_task-1)) * (X_task.T @ X_task - m_task * np.outer(x_bar_task, x_bar_task))

if S_task is not None and x_bar_task is not None:
    print("  S_task =")
    print(S_task)

    # Automatic check: compare with the definition form
    if X_task_tilde is not None:
        S_check = (1/(m_task-1)) * X_task_tilde.T @ X_task_tilde
        err = np.abs(S_task - S_check).max()
        print(f"\n  Check (difference from the definition form): {err:.2e}  {'✓' if err < 1e-8 else '✗'}")
else:
    print("  [Complete steps 1 and 2 before filling in S_task]")

# ══════════════════════════════════════════════════════════
# 🖊 Step 4: interpret S_task
# ══════════════════════════════════════════════════════════
print_step(4, "Interpret S_task: variances, standard deviations, correlation coefficients")

if S_task is not None:
    print("\n  Variance and standard deviation of each measure:")
    for j, name in enumerate(feature_names_task):
        var_j = S_task[j, j]
        std_j = np.sqrt(var_j)
        print(f"    {name}: Var = {var_j:.2f}, σ = {std_j:.2f}")

    print("\n  Pearson correlation coefficients r between the measures:")
    for i in range(n_task):
        for j in range(i+1, n_task):
            r_ij = S_task[i, j] / np.sqrt(S_task[i, i] * S_task[j, j])
            if abs(r_ij) > 0.7:
                interp = " (strong)"
            elif abs(r_ij) > 0.3:
                interp = " (moderate)"
            else:
                interp = " (weak)"
            sign_str = "positive correlation" if r_ij > 0 else "negative correlation"
            print(f"    r({feature_names_task[i]}, {feature_names_task[j]}) = {r_ij:+.4f}  → {sign_str}{interp}")
else:
    print("  [Complete step 3 first]")

# ══════════════════════════════════════════════════════════
# 🖊 Step 5: compute tr(S_task) and ||X̃||_F², and check the relation between them
# ══════════════════════════════════════════════════════════
print_step(5, "Check: the relation between tr(S_task) and ||X̃_task||_F²")

if S_task is not None and X_task_tilde is not None:
    # 🖊 Student fill-in:
    trace_S_task  = None  # change to: np.trace(S_task)
    frob_sq_task  = None  # change to: np.sum(X_task_tilde**2)

    if trace_S_task is not None and frob_sq_task is not None:
        lhs = (m_task - 1) * trace_S_task
        compare_print(
            "(m-1)*tr(S) vs ||X̃||_F²",
            f"{lhs:.4f}  vs  {frob_sq_task:.4f}",
            "the two should be equal"
        )
        print(f"\n  Total variance tr(S) = {trace_S_task:.4f}")
        print(f"  Average variance per feature = tr(S)/n = {trace_S_task/n_task:.4f}")
        print(f"  Estimated average standard deviation ≈ {np.sqrt(trace_S_task/n_task):.4f}")
    else:
        print("  [Fill in trace_S_task and frob_sq_task]")
else:
    print("  [Complete steps 3 and 2 first]")

# ============================================================
print_header("E | Visualization: the covariance structure of three-dimensional data")
# ============================================================

if 'S_task' in dir() and S_task is not None:

    # --- Compute the correlation matrix (derived from S_task) ---
    D_std = np.diag(1.0 / np.sqrt(np.diag(S_task)))
    R_task = D_std @ S_task @ D_std   # correlation matrix

    fig = make_subplots(
        rows=1, cols=2,
        subplot_titles=["Covariance matrix S", "Correlation matrix R"],
        horizontal_spacing=0.15
    )

    # Heat map of the covariance matrix
    fig.add_trace(go.Heatmap(
        z=S_task,
        colorscale=[[0, C_BG], [0.5, C_AUX], [1, C_V1]],
        text=[[f"{v:.1f}" for v in row] for row in S_task],
        texttemplate="%{text}", textfont=dict(size=13),
        showscale=True, colorbar=dict(x=0.44, title="Cov"),
        xgap=3, ygap=3
    ), row=1, col=1)

    # Heat map of the correlation matrix (fixed from -1 to 1)
    fig.add_trace(go.Heatmap(
        z=R_task,
        colorscale=[[0, C_WARN], [0.5, C_BG], [1, C_V1]],
        zmin=-1, zmax=1,
        text=[[f"{v:+.3f}" for v in row] for row in R_task],
        texttemplate="%{text}", textfont=dict(size=13),
        showscale=True, colorbar=dict(x=1.01, title="r"),
        xgap=3, ygap=3
    ), row=1, col=2)

    short = ["Reaction", "Grip", "Heart rate"]
    for col_idx in [1, 2]:
        fig.update_xaxes(
            tickvals=[0, 1, 2], ticktext=short, row=1, col=col_idx
        )
        fig.update_yaxes(
            tickvals=[0, 1, 2], ticktext=short,
            autorange='reversed', row=1, col=col_idx
        )

    fig.update_layout(
        title=dict(
            text="Three-dimensional physiological data: covariance matrix vs. correlation matrix",
            font=dict(size=14, color=C_AXIS)
        ),
        plot_bgcolor=C_BG, paper_bgcolor=C_BG,
        height=400, width=820,
        font=dict(family="Microsoft JhengHei, SimHei, sans-serif")
    )
    fig.show()

    # --- Scatter matrix (in place of a pairplot) ---
    print("\n  Drawing pairwise scatter plots of the features (to see the linear relations)...")

    fig2 = make_subplots(
        rows=n_task, cols=n_task,
        subplot_titles=[
            f"{short[i]} vs {short[j]}" if i != j else short[i]
            for i in range(n_task) for j in range(n_task)
        ],
        horizontal_spacing=0.08, vertical_spacing=0.12
    )

    pair_colors = [C_V1, C_V2, C_T1]
    for i in range(n_task):
        for j in range(n_task):
            if i == j:
                # diagonal: show the variance as text
                fig2.add_trace(go.Scatter(
                    x=[0], y=[0], mode='text',
                    text=[f"σ²={S_task[i,i]:.1f}<br>σ={np.sqrt(S_task[i,i]):.2f}"],
                    textfont=dict(size=12, color=C_V1),
                    showlegend=False
                ), row=i+1, col=j+1)
            else:
                if X_task_tilde is not None:
                    fig2.add_trace(go.Scatter(
                        x=X_task_tilde[:, j],
                        y=X_task_tilde[:, i],
                        mode='markers',
                        marker=dict(size=8, color=pair_colors[i], opacity=0.8),
                        showlegend=False,
                        hovertemplate=(
                            f"{short[j]}=%{{x:.1f}}<br>"
                            f"{short[i]}=%{{y:.1f}}"
                            f"<extra>r={R_task[i,j]:+.3f}</extra>"
                        )
                    ), row=i+1, col=j+1)

    fig2.update_layout(
        title=dict(
            text="Pairwise scatter plots of the centered features (hover to see the correlation coefficient)",
            font=dict(size=13, color=C_AXIS)
        ),
        plot_bgcolor=C_BG, paper_bgcolor=C_BG,
        height=520, width=620,
        font=dict(family="Microsoft JhengHei, SimHei, sans-serif")
    )
    fig2.update_xaxes(gridcolor=C_GRID, gridwidth=0.6, zeroline=True, zerolinecolor=C_AXIS)
    fig2.update_yaxes(gridcolor=C_GRID, gridwidth=0.6, zeroline=True, zerolinecolor=C_AXIS)
    fig2.show()

else:
    print("  [Complete step 3 and compute S_task before running this cell]")