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

Nonlinear exponential fitting by Nelder-Mead in A and log(tau).
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_nonlinear_fit.py [--no-plot]
Requires NumPy; optional plots require Matplotlib. Nonlinear fitting also requires SciPy.
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()
    t = np.linspace(0,5,30)
    y = 2.0*np.exp(-t/1.5)+0.1*rng.standard_normal(t.size)
    sigma_y = np.full(t.size,0.1)  # None for unweighted fitting
    A0,tau0 = 1.0,1.0
    max_iter,max_evaluations = 2000,4000
    tol_parameters,tol_objective = 1e-8,1e-8
    A,tau,info = exponential_fit(t,y,sigma_y,A0,tau0,max_iter,max_evaluations,tol_parameters,tol_objective)
    print(f"A = {A:.12g}\ntau = {tau:.12g}")
    for key,value in info.items():
        if key != "residuals": print(f"{key}: {value}")
    print("Point estimates only; inspect residuals and try other initial guesses.")
    if make_plot: plot_fit(t,y,lambda tt:A*np.exp(-tt/tau),info["residuals"],"Exponential fit",sigma_y)


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 exponential_fit(t, y, sigma_y=None, A0=1.0, tau0=1.0, max_iter=2000,
                    max_evaluations=4000, tol_parameters=1e-8, tol_objective=1e-8):
    """Nelder-Mead in A,log(tau); fit original y values, not log(y).

    This example reports point estimates, not parameter uncertainties.
    """
    from scipy.optimize import minimize
    t, y = vector(t,"t"), vector(y,"y")
    if t.shape != y.shape or t.size < 3 or np.any(t < 0) or np.ptp(t) <= 0:
        raise ValueError("Require >=3 matching observations at distinct nonnegative times")
    positive(tau0,"tau0")
    if not np.isfinite(A0):
        raise ValueError("A0 must be finite")
    integer(max_iter,1,"max_iter"); integer(max_evaluations,3,"max_evaluations")
    positive(tol_parameters,"tol_parameters");positive(tol_objective,"tol_objective")
    scale = np.ones_like(t) if sigma_y is None else vector(sigma_y,"sigma_y")
    if scale.shape != t.shape or np.any(scale <= 0):
        raise ValueError("Supply one positive standard deviation per observation")
    def cost(q):
        with np.errstate(over="ignore", under="ignore", invalid="ignore", divide="ignore"):
            tau = np.exp(q[1]); r = (y-q[0]*np.exp(-t/tau))/scale
            value = float(r@r)
        return value if np.isfinite(value) and np.isfinite(tau) and tau > 0 else 1e300
    result = minimize(cost, [A0,np.log(tau0)], method="Nelder-Mead",
                      options=dict(maxiter=max_iter,maxfev=max_evaluations,
                                   xatol=tol_parameters,fatol=tol_objective))
    A, tau = float(result.x[0]), float(np.exp(result.x[1]))
    if result.fun >= 1e300 or not np.isfinite(tau) or tau <= 0:
        raise ValueError("No numerically usable fit was found")
    residual = y-A*np.exp(-t/tau); sse = float(residual@residual)
    return A, tau, dict(residuals=residual,sse=sse,objective=float(result.fun),
                        residual_std=np.sqrt(sse/(t.size-2)),dof=t.size-2,
                        iterations=result.nit,evaluations=result.nfev,
                        converged=bool(result.success),status=str(result.message))

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_fit(x, y, prediction, residual, title, sigma=None):
    import matplotlib.pyplot as plt
    xx=np.linspace(np.min(x),np.max(x),300);fig,ax=plt.subplots(2,1,sharex=True)
    ax[0].errorbar(x,y,yerr=sigma,fmt="ko",mfc="white",label="Data")
    ax[0].plot(xx,prediction(xx),"k-",label="Fit");ax[0].legend(loc="best")
    ax[0].set_title(title,fontweight="normal");ax[0].set_ylabel("y")
    ax[1].errorbar(x,residual,yerr=sigma,fmt="ko",mfc="white")
    ax[1].axhline(0,color="k",ls="--");ax[1].set_xlabel("x");ax[1].set_ylabel("y - fit")
    for a in ax:a.grid(True)
    fig.tight_layout();plt.show()


if __name__ == "__main__":
    main()