"""ch13_simulated_annealing.py
Companion code for Chapter 13: Optimization and Data Fitting.

Bounded scalar simulated annealing.
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_simulated_annealing.py [--no-plot]
Requires NumPy; optional plots require Matplotlib.
Fixed seeds reproduce runs with the same software; languages need not share sequences.
"""

import sys
import numpy as np


def main():
    make_plot = "--no-plot" not in sys.argv
    seed = 1
    rng = np.random.default_rng(seed)
    # Alternative: comment out the previous line and use fresh system entropy.
    # rng = np.random.default_rng()
    fun = lambda x: x**4-20*x*x+24*x+5
    a,b,x0 = -5.0,4.0,2.0
    T0,T_min,cooling = 100.0,1e-3,0.95
    moves_per_temp,step_size,max_levels = 100,0.5,1000
    x,fx,info = simulated_annealing(fun,a,b,x0,rng,T0,T_min,cooling,moves_per_temp,step_size,max_levels)
    report(x,fx,info)
    print("Repeat with other seeds; global optimality is not guaranteed.")
    if make_plot: plot_scalar(fun,a,b,x,fx,"Simulated annealing")


def vector(x, name="data"):
    if np.iscomplexobj(x):
        raise ValueError(name + " must be real")
    a = np.asarray(x, dtype=float)
    if a.ndim != 1 or a.size == 0 or not np.all(np.isfinite(a)):
        raise ValueError(name + " must be a nonempty finite vector")
    return a


def checked(fun, x):
    value = np.asarray(fun(x))
    if value.ndim != 0 or np.iscomplexobj(value) or not np.isfinite(value):
        raise ValueError("Objective must return a finite real scalar")
    return float(value)


def bounds(lower, upper):
    lo, hi = vector(lower, "lower"), vector(upper, "upper")
    if lo.shape != hi.shape or not np.all(np.isfinite(hi - lo)) or np.any(hi <= lo):
        raise ValueError("Require matching finite bounds with upper > lower")
    return lo, hi


def integer(value, minimum, name):
    if not np.isscalar(value) or not np.isfinite(value) or int(value) != value or value < minimum:
        raise ValueError(name + " must be an integer >= " + str(minimum))
    return int(value)


def positive(value, name):
    if not np.isscalar(value) or not np.isfinite(value) or value <= 0:
        raise ValueError(name + " must be positive and finite")
    return float(value)

def simulated_annealing(fun, a, b, x0, rng, T0=100.0, T_min=1e-3,
                        cooling=0.95, moves_per_temp=100, step_size=0.5, max_levels=1000):
    """Symmetric uniform proposals; reject out-of-bounds moves without clipping."""
    bounds([a], [b]); positive(T0, "T0"); positive(T_min, "T_min")
    positive(step_size, "step_size")
    if not np.isfinite(x0) or not a <= x0 <= b or not 0 < T_min < T0 or not 0 < cooling < 1:
        raise ValueError("Invalid start or cooling settings")
    moves_per_temp = integer(moves_per_temp, 1, "moves_per_temp")
    max_levels = integer(max_levels, 1, "max_levels")
    x = float(x0); fx = checked(fun, x); best, fb = x, fx
    T = T0; levels = trials = accepted = outside = 0; calls = 1
    while T > T_min and levels < max_levels:
        for _ in range(moves_per_temp):
            trials += 1; xt = x + step_size*rng.uniform(-1, 1)
            if not a <= xt <= b:
                outside += 1; continue
            ft = checked(fun, xt); calls += 1; df = ft-fx
            if df <= 0 or np.log(max(rng.random(), np.finfo(float).tiny)) < -df/T:
                x, fx = xt, ft; accepted += 1
                if fx < fb:
                    best, fb = x, fx
        levels += 1; T *= cooling
    status = "Cooling schedule completed." if T <= T_min else "Temperature-level limit reached."
    return best, fb, dict(iterations=levels, trials=trials, accepted=accepted,
                          outside=outside, evaluations=calls, final_temperature=T, status=status)

def report(p, value, info):
    for j,x in enumerate(np.atleast_1d(p),1):
        print(f"p({j}) = {x:.12g}")
    print(f"Objective value: {value:.12g}")
    for key,value in info.items():
        print(f"{key}: {value}")

def plot_scalar(fun, a, b, x, fx, title):
    import matplotlib.pyplot as plt
    xx=np.linspace(a,b,400)
    plt.plot(xx,[fun(z) for z in xx],"k-",label="Objective")
    plt.plot(x,fx,"ko",label="Best point");plt.xlabel("x");plt.ylabel("f(x)")
    plt.title(title,fontweight="normal");plt.xlim(a,b);plt.legend(loc="best");plt.grid(True);plt.tight_layout();plt.show()


if __name__ == "__main__":
    main()