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

Polynomial model comparison using independent validation data.
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_polynomial_models.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 = 7
    rng = np.random.default_rng(seed)
    # Alternative: comment out the previous line and use fresh system entropy.
    # rng = np.random.default_rng()
    true_function = lambda x: np.sin(2*np.pi*x)+0.5*x
    noise_sigma = 0.18
    x = np.linspace(0,1,18); y = true_function(x)+noise_sigma*rng.standard_normal(x.size)
    xv = np.linspace(0,1,101); yv = true_function(xv)+noise_sigma*rng.standard_normal(xv.size)
    degrees = [1,5,14]
    # For measured data, use separate training and validation measurements.
    # Reserve fresh test data after selecting a model; never fit validation data.
    models=[]
    print("Degree       Training RMSE    Validation RMSE")
    for degree in degrees:
        model=polynomial_fit(x,y,degree);models.append(model)
        train=np.sqrt(np.mean((y-polynomial_predict(model,x))**2))
        valid=np.sqrt(np.mean((yv-polynomial_predict(model,xv))**2))
        print(f"{degree:6d}       {train:.8g}       {valid:.8g}")
    if make_plot:
        import matplotlib.pyplot as plt
        fig,ax=plt.subplots(2,1,sharex=True);xx=np.linspace(0,1,600)
        ax[0].plot(x,y,"ko",mfc="white",label="Training data")
        styles=["-","--",":"]
        for model,degree,style in zip(models,degrees,styles):
            ax[0].plot(xx,polynomial_predict(model,xx),color="k",ls=style,label=f"Degree {degree}")
            ax[1].plot(x,y-polynomial_predict(model,x),color="k",ls=style,marker="o",label=f"Degree {degree}")
        ax[0].set_ylim(-1.5,2);ax[0].set_title("Polynomial models",fontweight="normal")
        ax[0].set_ylabel("y");ax[0].legend(loc="best");ax[1].axhline(0,color="0.5")
        ax[1].set_xlabel("x");ax[1].set_ylabel("Training residual")
        for a in ax:a.grid(True);a.set_xlim(0,1)
        fig.tight_layout();plt.show()


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 polynomial_fit(x, y, degree):
    """Unweighted least squares in centered/scaled monomials, solved by SVD."""
    x,y = vector(x,"x"),vector(y,"y"); degree = integer(degree,0,"degree")
    if y.shape != x.shape or x.size <= degree or np.unique(x).size <= degree or x.size < 2:
        raise ValueError("Require more distinct x values than the polynomial degree")
    center, scale = float(np.mean(x)),float(np.std(x,ddof=1))
    positive(scale,"x scale")
    A = np.polynomial.polynomial.polyvander((x-center)/scale, degree)
    coeff,_,rank,_ = np.linalg.lstsq(A,y,rcond=None)
    if rank != degree+1:
        raise ValueError("Rank-deficient polynomial design; lower the degree")
    return dict(coefficients=coeff,center=center,scale=scale)


def polynomial_predict(model, x):
    return np.polynomial.polynomial.polyval((np.asarray(x)-model["center"])/model["scale"], model["coefficients"])

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}")


if __name__ == "__main__":
    main()