
"""PKTron Algorithms Science — Optimization Algorithms."""
import numpy as np
_R = np.float32

def gradient_descent(f, grad_f, x0, lr=0.01, n_iter=500, tol=1e-6):
    """
    Vanilla gradient descent.
    Returns dict: x_opt, f_opt, history, n_iter_run.
    """
    x    = np.asarray(x0, dtype=_R).copy()
    hist = []
    for i in range(n_iter):
        g = np.asarray(grad_f(x), dtype=_R)
        x = x - _R(lr)*g
        fv = float(f(x))
        hist.append(fv)
        if np.linalg.norm(g) < tol:
            break
    return {"x_opt":x,"f_opt":float(f(x)),"history":hist,"n_iter_run":i+1}

def adam_optimizer(f, grad_f, x0, lr=0.001, beta1=0.9, beta2=0.999,
                   eps=1e-8, n_iter=1000, tol=1e-7):
    """Adam optimizer. Returns dict same as gradient_descent."""
    x  = np.asarray(x0, dtype=_R).copy()
    m  = np.zeros_like(x)
    v  = np.zeros_like(x)
    hist=[]
    for i in range(1, n_iter+1):
        g = np.asarray(grad_f(x), dtype=_R)
        m = _R(beta1)*m + _R(1-beta1)*g
        v = _R(beta2)*v + _R(1-beta2)*g**2
        m_hat = m/(_R(1-beta1**i))
        v_hat = v/(_R(1-beta2**i))
        x = x - _R(lr)*m_hat/(np.sqrt(v_hat)+_R(eps))
        fv=float(f(x)); hist.append(fv)
        if np.linalg.norm(g)<tol: break
    return {"x_opt":x,"f_opt":float(f(x)),"history":hist,"n_iter_run":i}

def simulated_annealing(f, x0, T_start=1.0, T_end=1e-4, n_steps=5000,
                         step_size=0.1, seed=0):
    """Simulated annealing minimiser. Returns dict."""
    rng  = np.random.default_rng(seed)
    x    = np.asarray(x0, dtype=_R).copy()
    fx   = float(f(x))
    best_x, best_f = x.copy(), fx
    hist = [fx]
    for i in range(n_steps):
        T     = _R(T_start)*((_R(T_end)/_R(T_start))**(_R(i)/n_steps))
        x_new = x + _R(step_size)*rng.standard_normal(x.shape).astype(_R)
        fn    = float(f(x_new))
        delta = fn - fx
        if delta < 0 or rng.random() < float(np.exp(-_R(delta)/T)):
            x=x_new; fx=fn
        if fx < best_f: best_x=x.copy(); best_f=fx
        hist.append(fx)
    return {"x_opt":best_x,"f_opt":best_f,"history":hist}

def nelder_mead_simple(f, x0, max_iter=500, tol=1e-5, seed=0):
    """Wrap scipy Nelder-Mead if available, else gradient descent."""
    try:
        from scipy.optimize import minimize as sp_min
        res = sp_min(f, x0, method="Nelder-Mead",
                     options={"maxiter":max_iter,"xatol":tol,"fatol":tol})
        return {"x_opt":res.x,"f_opt":float(res.fun),"n_iter_run":res.nit}
    except ImportError:
        return gradient_descent(f, lambda x: x*0, x0, n_iter=max_iter)

def genetic_algorithm_float(f, bounds, pop_size=50, n_gen=100,
                              mutation_std=0.1, seed=0):
    """
    Float-vector genetic algorithm minimiser.
    bounds: list of (min,max) per dimension.
    """
    rng = np.random.default_rng(seed)
    dim = len(bounds)
    lo  = np.array([b[0] for b in bounds], dtype=_R)
    hi  = np.array([b[1] for b in bounds], dtype=_R)
    pop = rng.uniform(lo, hi, (pop_size, dim)).astype(_R)
    hist=[]
    for gen in range(n_gen):
        scores = np.array([f(ind) for ind in pop], dtype=_R)
        hist.append(float(scores.min()))
        idx    = np.argsort(scores)
        pop    = pop[idx]
        n_elite= max(2, pop_size//5)
        new_pop = list(pop[:n_elite])
        while len(new_pop) < pop_size:
            p1 = pop[rng.integers(0,n_elite)]
            p2 = pop[rng.integers(0,n_elite)]
            child = _R(0.5)*(p1+p2)+_R(mutation_std)*rng.standard_normal(dim).astype(_R)
            child = np.clip(child, lo, hi)
            new_pop.append(child)
        pop = np.array(new_pop[:pop_size], dtype=_R)
    scores  = np.array([f(ind) for ind in pop], dtype=_R)
    best_i  = int(np.argmin(scores))
    return {"x_opt":pop[best_i],"f_opt":float(scores[best_i]),"history":hist}
