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

Real-valued genetic algorithm with elitism.
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_genetic_algorithm.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 = 42
    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.cos(3*p[1])+(p[0]-1)**2+(p[1]+2)**2
    lower,upper = np.array([-3.0,-3.0]),np.array([3.0,3.0])
    population_size,max_generations = 100,200
    crossover_probability,mutation_probability = 0.9,0.3
    mutation_sigma,tournament_size = 0.04*(upper-lower),3
    p,fx,info = genetic_search(fun,lower,upper,rng,population_size,max_generations,crossover_probability,mutation_probability,mutation_sigma,tournament_size)
    report(p,fx,info)
    print("Repeat with other seeds; global optimality is not guaranteed.")
    if make_plot: plot_contour(fun,[-1,3],[-3,0],p,"Genetic algorithm")


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 genetic_search(fun, lower, upper, rng, population_size=100, max_generations=200,
                   crossover_probability=0.9, mutation_probability=0.3,
                   mutation_sigma=None, tournament_size=3):
    """Real-valued GA: tournaments, uniform crossover, Gaussian mutation, one elite."""
    lo, hi = bounds(lower, upper); d = lo.size
    n = integer(population_size, 2, "population_size")
    K = integer(max_generations, 1, "max_generations")
    nt = integer(tournament_size, 1, "tournament_size")
    sigma = vector(mutation_sigma if mutation_sigma is not None else 0.04*(hi-lo), "mutation_sigma")
    if sigma.shape != lo.shape or np.any(sigma <= 0) or not 0 <= crossover_probability <= 1 or not 0 <= mutation_probability <= 1:
        raise ValueError("Invalid mutation widths or probabilities")
    pop = rng.uniform(lo, hi, (n,d)); values = np.array([checked(fun,p) for p in pop])
    def parent():
        idx = rng.integers(n, size=nt)
        return pop[idx[np.argmin(values[idx])]].copy()
    for _ in range(K):
        j = np.argmin(values); nxt, nv = [pop[j].copy()], [values[j]]
        while len(nxt) < n:
            a, b = parent(), parent(); children = [a.copy(), b.copy()]
            if rng.random() < crossover_probability:
                mask = rng.random(d) < 0.5
                children[0][mask], children[1][mask] = b[mask], a[mask]
            for child in children:
                if len(nxt) == n:
                    break
                if rng.random() < mutation_probability:
                    child += sigma*rng.standard_normal(d)
                child = np.clip(child, lo, hi)
                nxt.append(child); nv.append(checked(fun, child))
        pop, values = np.array(nxt), np.array(nv)
    j = np.argmin(values)
    return pop[j].copy(), float(values[j]), dict(iterations=K, evaluations=n+K*(n-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()