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

Current-to-best/1/bin differential evolution.
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_differential_evolution.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 = 3
    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.sum((p-0.5)**2+10*(1-np.cos(2*np.pi*(p-0.5))))
    lower,upper = [-4.5,-4.5],[5.5,5.5]
    population_size,max_generations,F,CR = 50,400,0.9,0.9
    p,fx,info = differential_evolution(fun,lower,upper,rng,population_size,max_generations,F,CR)
    report(p,fx,info)
    print("Repeat with other seeds; global optimality is not guaranteed.")
    if make_plot: plot_contour(fun,[-1,2],[-1,2],p,"Differential evolution")


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 differential_evolution(fun, lower, upper, rng, population_size=50,
                           max_generations=400, F=0.9, CR=0.9):
    """Frozen-generation current-to-best/1/bin DE, with uniform bound repair.

    Two distinct donors exclude the target; one mutant coordinate is compulsory.
    No local polishing is performed. The evaluation count includes initialization.
    """
    lo, hi = bounds(lower, upper); d = lo.size
    n = integer(population_size, 4, "population_size")
    K = integer(max_generations, 1, "max_generations")
    if not 0 < F <= 2 or not 0 <= CR <= 1:
        raise ValueError("Require 0 < F <= 2 and 0 <= CR <= 1")
    pop = rng.uniform(lo, hi, (n,d)); values = np.array([checked(fun,p) for p in pop])
    for _ in range(K):
        best = pop[np.argmin(values)].copy(); nxt, nv = pop.copy(), values.copy()
        for i in range(n):
            donors = rng.choice(np.delete(np.arange(n), i), 2, replace=False)
            mutant = pop[i] + F*(best-pop[i]) + F*(pop[donors[0]]-pop[donors[1]])
            mask = rng.random(d) < CR; mask[rng.integers(d)] = True
            trial = np.where(mask, mutant, pop[i]); bad = (trial < lo) | (trial > hi) | ~np.isfinite(trial)
            trial[bad] = rng.uniform(lo[bad], hi[bad])
            value = checked(fun, trial)
            if value <= values[i]:
                nxt[i], nv[i] = trial, value
        pop, values = nxt, nv
    j = np.argmin(values)
    return pop[j].copy(), float(values[j]), 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()