import numpy as np
import scipy.stats as st
import matplotlib.pyplot as plt

def compute_thresholds_normal(n, x_min=-6, x_max=6, dx=0.01):
    """
    Compute optimal stopping thresholds t_i for selecting the maximum
    of n i.i.d. N(0,1) draws using dynamic programming.
    """
    xs = np.arange(x_min, x_max+dx, dx)
    f = st.norm.pdf(xs)
    F = st.norm.cdf(xs)

    # V_{n+1}(x) = 0 for all x
    V_next = np.zeros_like(xs)

    thresholds = np.full(n+1, np.nan)
    V_store = [None]*(n+1)
    V_store[n] = V_next.copy()  # store V_{n+1}

    for i in range(n, 0, -1):
        power = n - i
        p_y = F**power if power > 0 else np.ones_like(xs)

        best_if_y = np.maximum(p_y, V_next)

        cum_f = np.cumsum(f)*dx
        cum_bestf = np.cumsum(best_if_y * f)*dx
        total_best = cum_bestf[-1]

        V_i = V_next * cum_f + (total_best - cum_bestf)

        V_store[i-1] = V_i.copy()
        V_next = V_i

        # Find threshold t_i: stop if p_y >= V_{i+1}(y)
        V_ip1 = V_store[i]
        mask = p_y >= V_ip1
        if np.any(mask):
            thresholds[i] = xs[np.where(mask)[0][0]]

    # Approximate success probability from start
    V1_start = V_store[0][0]

    return thresholds[1:], V1_start

# Example usage
n = 20
thresholds, V1_start = compute_thresholds_normal(n)

print(f"Initial success probability ≈ {V1_start:.3f}")
for i, t in enumerate(thresholds, 1):
    print(f"t_{i:2d} ≈ {t:.2f}")

plt.plot(range(1, n+1), thresholds, marker='o')
plt.xlabel("Stage i")
plt.ylabel("Threshold t_i")
plt.title(f"DP stop-thresholds (Standard Normal, n={n})")
plt.grid(True)
plt.show()