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

Bounded particle swarm optimization.
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_particle_swarm.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 = 2
    rng = np.random.default_rng(seed)
    # Alternative: comment out the previous line and use fresh system entropy.
    # rng = np.random.default_rng()
    fun = lambda p: np.sin(3*p[0])*np.sin(3*p[1])+0.1*np.dot(p,p)
    lower,upper = [-3.0,-3.0],[3.0,3.0]
    swarm_size,max_iter = 25,150
    inertia,c_personal,c_global = 0.72,1.4,1.4
    vmax = [0.5,0.5]
    p,fx,info = particle_swarm(fun,lower,upper,rng,swarm_size,max_iter,inertia,c_personal,c_global,vmax)
    report(p,fx,info)
    print("Repeat with other seeds; global optimality is not guaranteed.")
    if make_plot: plot_contour(fun,[-3,3],[-3,3],p,"Particle swarm optimization")


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 particle_swarm(fun, lower, upper, rng, swarm_size=25, max_iter=150,
                   inertia=0.72, c_personal=1.4, c_global=1.4, vmax=None):
    """Synchronous PSO; reflect boundary overshoots and halve reflected velocities."""
    lo, hi = bounds(lower, upper); width = hi-lo; d = lo.size
    n = integer(swarm_size, 2, "swarm_size"); K = integer(max_iter, 1, "max_iter")
    vm = vector(vmax if vmax is not None else 0.1*width, "vmax")
    if vm.shape != lo.shape or np.any(vm <= 0) or np.any(vm > width):
        raise ValueError("Require 0 < vmax <= upper-lower in each coordinate")
    if not 0 <= inertia < 1 or not np.isfinite(c_personal+c_global) or min(c_personal,c_global) < 0:
        raise ValueError("Invalid PSO coefficients")
    X = rng.uniform(lo, hi, (n,d)); V = rng.uniform(-vm, vm, (n,d))
    personal = X.copy(); values = np.array([checked(fun,p) for p in X])
    best = personal[np.argmin(values)].copy()
    for _ in range(K):
        V = inertia*V + c_personal*rng.random((n,d))*(personal-X) + c_global*rng.random((n,d))*(best-X)
        V = np.clip(V, -vm, vm); X += V
        below, above = X < lo, X > hi
        X = np.where(below, 2*lo-X, X); X = np.where(above, 2*hi-X, X)
        V[below | above] *= -0.5; X = np.clip(X, lo, hi)
        for i in range(n):
            value = checked(fun, X[i])
            if value < values[i]:
                personal[i], values[i] = X[i].copy(), value
        best = personal[np.argmin(values)].copy()
    return best, float(np.min(values)), dict(iterations=K, evaluations=n*(K+1), status="Search budget completed.")

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_contour(fun, lower, upper, p, title, start=None):
    import matplotlib.pyplot as plt
    X,Y=np.meshgrid(np.linspace(*lower,150),np.linspace(*upper,150))
    Z=np.array([fun(np.array([x,y])) for x,y in zip(X.ravel(),Y.ravel())]).reshape(X.shape)
    plt.contour(X,Y,Z,25,colors="0.6")
    if start is not None:
        plt.plot(*start,"ko",fillstyle="none",label="Start")
    plt.plot(*p,"k+",ms=10,label="Best point");plt.xlabel("p1");plt.ylabel("p2")
    plt.title(title,fontweight="normal");plt.xlim(lower);plt.ylim(upper);plt.gca().set_aspect("equal")
    plt.legend(loc="best");plt.grid(True);plt.tight_layout();plt.show()


if __name__ == "__main__":
    main()