"""Reproducible companions for /blogs/bayesian-optimization.

Run: python bayesian_optimization.py --output-dir bo-output
Dependencies: numpy, scipy, matplotlib, scikit-learn.
The toy function and fixed validation objective are deterministic; jitter is
numerical regularization, not simulated observation noise. No external data.
"""
from argparse import ArgumentParser
from pathlib import Path
from importlib.metadata import version
import json
import platform

import numpy as np
from scipy.linalg import cho_factor, cho_solve, solve_triangular
from scipy.special import ndtr


def kernel(x, z, length=0.18):
    """Unit-amplitude squared-exponential kernel on normalized inputs."""
    x, z = np.atleast_2d(x), np.atleast_2d(z)
    return np.exp(-0.5 * np.sum(((x[:, None] - z[None, :]) / length) ** 2, axis=2))


def posterior(x, y, candidates, length=0.18, noise_variance=0.0, jitter=1e-10):
    """Zero-mean GP: return latent mean/std, not noisy predictive std."""
    x, candidates = np.asarray(x), np.asarray(candidates)
    covariance = kernel(x, x, length) + (noise_variance + jitter) * np.eye(len(x))
    factor = cho_factor(covariance, lower=True)
    cross = kernel(x, candidates, length)
    mu = cross.T @ cho_solve(factor, y)
    projected = solve_triangular(factor[0], cross, lower=True)
    variance = 1.0 - np.sum(projected ** 2, axis=0)
    if np.min(variance) < -1e-7:
        raise ArithmeticError("Materially negative GP variance")
    return mu, np.sqrt(np.maximum(variance, 0.0))


def expected_improvement(mu, sd, best):
    """E[(F-best)+] for F ~ N(mu, sd**2), including sd=0."""
    mu, sd = np.broadcast_arrays(np.asarray(mu, float), np.asarray(sd, float))
    delta = mu - best
    z = np.divide(delta, sd, out=np.zeros_like(delta), where=sd > 0)
    density = np.exp(-0.5 * z * z) / np.sqrt(2 * np.pi)
    return np.where(sd > 0, delta * ndtr(z) + sd * density, np.maximum(delta, 0.0))


def objective(x):
    x = np.asarray(x)
    return 0.7 * np.sin(9 * x) + 0.3 * np.cos(17 * x) + 0.5 * x


def toy_example():
    grid = np.linspace(0, 1, 1001)[:, None]
    x = np.array([[0.05], [0.45], [0.90]])
    y = objective(x[:, 0])
    snapshots = []
    for _ in range(6):
        mu, sd = posterior(x, y, grid)
        ei = expected_improvement(mu, sd, y.max())
        # Never spend another noiseless evaluation on an already observed point.
        available = np.min(np.abs(grid - x[:, 0]), axis=1) > 1e-8
        index = np.argmax(np.where(available, ei, -np.inf))
        snapshots.append(dict(x=x.copy(), y=y.copy(), mu=mu, sd=sd, ei=ei, next=index))
        x = np.vstack([x, grid[index]])
        y = np.append(y, objective(grid[index, 0]))
    return grid, snapshots, x, y


def tuning_example(seed=19, budget=20, initial=5, pool_size=800):
    """Two optimizers, equal objective-call budgets, one untouched test set.

    Candidate pool is finite; random search samples uniformly without replacement.
    CV splits and data are fixed, so repeated evaluations have the same score.
    This surrogate is a working model, not a GP-prior theorem experiment.
    """
    from sklearn.datasets import make_moons
    from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    from sklearn.svm import SVC

    features, labels = make_moons(n_samples=500, noise=0.25, random_state=8)
    train_x, test_x, train_y, test_y = train_test_split(
        features, labels, test_size=0.25, stratify=labels, random_state=11)
    folds = list(StratifiedKFold(3, shuffle=True, random_state=23).split(train_x, train_y))
    rng = np.random.default_rng(seed)
    pool = rng.uniform(size=(pool_size, 2))
    permutation = rng.permutation(pool_size)
    start = permutation[:initial].tolist()

    def model(point):
        c, gamma = 10.0 ** (-3 + 6 * point)
        return make_pipeline(StandardScaler(), SVC(C=c, gamma=gamma))

    def score(point):
        # Pipeline fitting happens separately inside every training fold.
        return cross_val_score(model(point), train_x, train_y, cv=folds,
                               scoring="accuracy", n_jobs=1).mean()

    initial_scores = [score(pool[i]) for i in start]
    records = {}
    for method in ("EI", "Random search"):
        chosen, values = start.copy(), initial_scores.copy()
        while len(chosen) < budget:
            if method == "EI":
                values_array = np.asarray(values)
                center = values_array.mean()
                scale = max(values_array.std(), 1e-3)
                mu, sd = posterior(pool[chosen], (values_array - center) / scale,
                                   pool, length=0.25)
                acquisition = expected_improvement(mu, sd, (max(values) - center) / scale)
                acquisition[chosen] = -np.inf
                index = int(acquisition.argmax())
            else:
                index = int(permutation[len(chosen)])
            chosen.append(index)
            values.append(float(score(pool[index])))
        winner = int(np.argmax(values))
        final_model = model(pool[chosen[winner]]).fit(train_x, train_y)
        # Test labels are read only after this method's CV-based selection ends.
        test_accuracy = float(final_model.score(test_x, test_y))
        records[method] = dict(indices=chosen, scores=values,
                               best_cv=float(max(values)), test_accuracy=test_accuracy,
                               parameters=dict(zip(("C", "gamma"), (10.0 ** (-3 + 6 * pool[chosen[winner]])).tolist())))
    return records


def self_check():
    """Numerical checks of EI, GP conditioning and log-det factorization."""
    from scipy.integrate import quad
    from scipy.stats import norm
    for mu, sd, best in [(0.3, 0.7, 0.5), (-2., 0.2, 0.), (2., 0.1, 0.)]:
        integral = quad(lambda z: max(mu + sd * z - best, 0) * norm.pdf(z),
                        -12, 12, points=[(best - mu) / sd] if abs((best - mu) / sd) < 12 else None)[0]
        np.testing.assert_allclose(expected_improvement(mu, sd, best), integral, atol=1e-9)
    np.testing.assert_allclose(expected_improvement(np.array([2., -2.]), np.zeros(2), 0), [2., 0.])
    x = np.array([[0.1], [0.4], [0.8]])
    y = np.array([0.2, -0.4, 0.7])
    points = np.array([[0.2], [0.65]])
    tau2 = 0.1
    mu, sd = posterior(x, y, points, noise_variance=tau2, jitter=0)
    a = kernel(x, x) + tau2 * np.eye(len(x))
    cross = kernel(x, points)
    np.testing.assert_allclose(mu, cross.T @ np.linalg.solve(a, y), atol=1e-12)
    np.testing.assert_allclose(sd ** 2, np.diag(kernel(points, points) - cross.T @ np.linalg.solve(a, cross)), atol=1e-12)
    log_increment = 0.
    for i, point in enumerate(x):
        variance = 1. if i == 0 else posterior(x[:i], y[:i], point[None], noise_variance=tau2, jitter=0)[1][0] ** 2
        log_increment += np.log1p(variance / tau2)
    np.testing.assert_allclose(log_increment, np.linalg.slogdet(np.eye(len(x)) + kernel(x, x) / tau2)[1], atol=1e-12)
    print("EI quadrature, zero-variance limit, GP moments and log-det checks passed.")


def save_plots(output_dir, grid, snapshots, records):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    palette = dict(blue="#315f9a", gold="#b47723", ink="#273446", grid="#dce2e9")
    plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 12,
                         "axes.spines.top": False, "axes.spines.right": False,
                         "axes.labelcolor": palette["ink"], "text.color": palette["ink"],
                         "svg.fonttype": "none", "savefig.facecolor": "white"})
    fig, axes = plt.subplots(4, 1, figsize=(8.5, 10), layout="constrained",
                             gridspec_kw={"height_ratios": [2, 1, 2, 1]})
    for row, snapshot in enumerate((snapshots[0], snapshots[3])):
        ax, acquisition = axes[2 * row:2 * row + 2]
        xx, next_x = grid[:, 0], grid[snapshot["next"], 0]
        ax.plot(xx, objective(xx), "--", color=palette["ink"], label="True function (for illustration)")
        ax.plot(xx, snapshot["mu"], color=palette["blue"], label="Posterior mean")
        ax.fill_between(xx, snapshot["mu"] - 1.96 * snapshot["sd"],
                        snapshot["mu"] + 1.96 * snapshot["sd"], color=palette["blue"], alpha=0.13,
                        label="95% pointwise posterior interval")
        ax.scatter(snapshot["x"][:, 0], snapshot["y"], color=palette["ink"], zorder=5, s=32, label="Evaluated")
        ax.set(ylabel="Function value", title=f"After {len(snapshot['x'])} evaluations", xlim=(0, 1), ylim=(-2.3, 2.7))
        if row == 0:
            ax.legend(fontsize=9, loc="upper left", ncol=2)
        acquisition.plot(xx, snapshot["ei"], color=palette["gold"])
        acquisition.scatter(next_x, snapshot["ei"][snapshot["next"]], marker="D", color=palette["gold"], zorder=5)
        acquisition.annotate(f"Next query: {next_x:.3f}", (next_x, snapshot["ei"][snapshot["next"]]),
                             xytext=(10 if next_x < .7 else -130, 10), textcoords="offset points", fontsize=10)
        acquisition.set(xlim=(0, 1), ylim=(0, max(snapshot["ei"]) * 1.55), ylabel="Expected\nimprovement", xlabel="Input x")
        for axis in (ax, acquisition):
            axis.grid(axis="y", color=palette["grid"], linewidth=.6)
    fig.savefig(output_dir / "bo-posterior-ei.svg")
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(8.5, 4.3), layout="constrained")
    for method, color, linestyle in [("EI", palette["blue"], "-"), ("Random search", palette["gold"], "--")]:
        scores = np.maximum.accumulate(records[method]["scores"]) * 100
        ax.step(np.arange(1, len(scores) + 1), scores, where="post", label=method, color=color, linestyle=linestyle)
    ax.axvline(5, color=palette["ink"], linewidth=.8, linestyle=":")
    ax.text(5.3, .04, "Shared initial design", transform=ax.get_xaxis_transform(), fontsize=10)
    ax.set(xlabel="Objective evaluations (3-fold CV each)", ylabel="Best CV accuracy (%)", xlim=(1, 20),
           xticks=[1, 5, 10, 15, 20], ylim=(45, 100))
    ax.grid(axis="y", color=palette["grid"], linewidth=.6)
    ax.legend(loc="lower right")
    fig.savefig(output_dir / "bo-tuning-budget.svg")
    plt.close(fig)


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path("bo-output"))
    parser.add_argument("--check-only", action="store_true")
    args = parser.parse_args()
    self_check()
    if args.check_only:
        return
    args.output_dir.mkdir(parents=True, exist_ok=True)
    grid, snapshots, x, y = toy_example()
    records = tuning_example()
    save_plots(args.output_dir, grid, snapshots, records)
    results = dict(metadata=dict(python=platform.python_version(),
                                packages={name: version(name) for name in ("numpy", "scipy", "scikit-learn", "matplotlib")},
                                data_seed=8, split_seed=11, cv_seed=23, search_seed=19,
                                candidate_pool=800, budget_per_method=20, common_initial=5,
                                train_size=375, test_size=125, cv_folds=3),
                   toy=dict(queries=x[:, 0].tolist(), values=y.tolist()), tuning=records)
    (args.output_dir / "bo-results.json").write_text(json.dumps(results, indent=2) + "\n")
    print(json.dumps({name: {k: v for k, v in record.items() if k not in ("indices", "scores")}
                      for name, record in records.items()}, indent=2))


if __name__ == "__main__":
    main()