#!/usr/bin/env python3
import argparse
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import linprog


def make_data(n=41, tnoise='z2', knoise=0.8, seed=0):
    if seed < 0: seed = None
    rng = np.random.default_rng(seed)
    x = np.linspace(0, 1, n)
    y_true = np.exp(2.2 * x) - 1.0

    if tnoise == 'n':
# normal distribution
        z = rng.normal(size=n)
        noise = knoise * z
    elif tnoise == 'z2':
# extneded, asymmetric distribution
        z = rng.normal(size=n)
        noise = knoise * (z**2 - 1.0) / np.sqrt(2.0)
    elif tnoise == 't':
# t distribution, extended tail, symmetric
        z = rng.standard_t(df=3, size=n)
# t分布(df=3)の標準偏差を1にそろえる
        noise = knoise * z / np.sqrt(3.0)

    y = y_true + noise
    
    return x, y, y_true


def design_matrix(x, degree):
    return np.vstack([x**k for k in range(degree + 1)]).T


def predict(x, coef):
    X = design_matrix(x, len(coef) - 1)
    return X @ coef


def least_squares_fit(x, y, degree):
    X = design_matrix(x, degree)
    coef, *_ = np.linalg.lstsq(X, y, rcond=None)
    return coef


def minimax_lp_fit(x, y, degree):
    X = design_matrix(x, degree)
    n, p = X.shape

    # unknown = [a0, a1, ..., ap, t]
    c = np.zeros(p + 1)
    c[-1] = 1.0

    A = []
    b = []

    # X a - y <= t  ->  X a - t <= y
    A.append(np.hstack([X, -np.ones((n, 1))]))
    b.append(y)

    # -(X a - y) <= t  ->  -X a - t <= -y
    A.append(np.hstack([-X, -np.ones((n, 1))]))
    b.append(-y)

    A_ub = np.vstack(A)
    b_ub = np.concatenate(b)

    bounds = [(None, None)] * p + [(0, None)]

    res = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method="highs")
    if not res.success:
        raise RuntimeError(res.message)

    coef = res.x[:p]
    max_error = res.x[-1]
    return coef, max_error


def solve_alternating_system(x, y, degree, active):
    active = np.array(active, dtype=int)
    xa = x[active]
    ya = y[active]

    X = design_matrix(xa, degree)
    signs = np.array([(-1) ** i for i in range(len(active))], dtype=float)

    # X a - y = signs * E
    # X a - signs * E = y
    A = np.column_stack([X, -signs])

    sol = np.linalg.solve(A, ya)
    coef = sol[:-1]
    E = sol[-1]
    return coef, abs(E)


def minimax_remez_fit(x, y, degree, maxiter=100, tol=1e-10):
    n = len(x)
    m = degree + 2

    active = np.linspace(0, n - 1, m, dtype=int)
    best_coef = None
    best_err = np.inf

    for _ in range(maxiter):
        coef, _ = solve_alternating_system(x, y, degree, active)

        r = predict(x, coef) - y
        err = np.max(np.abs(r))

        if err < best_err:
            best_err = err
            best_coef = coef.copy()

        candidates = set(np.argsort(np.abs(r))[-m * 3:])
        candidates.update(active.tolist())

        improved = False
        best_trial = (err, active.copy(), coef.copy())

        for j in range(m):
            for cand in candidates:
                trial = active.copy()
                trial[j] = cand
                trial = np.unique(trial)

                if len(trial) != m:
                    continue

                trial.sort()

                try:
                    c_trial, _ = solve_alternating_system(x, y, degree, trial)
                except np.linalg.LinAlgError:
                    continue

                r_trial = predict(x, c_trial) - y
                e_trial = np.max(np.abs(r_trial))

                if e_trial + tol < best_trial[0]:
                    best_trial = (e_trial, trial.copy(), c_trial.copy())
                    improved = True

        active = best_trial[1]
        best_coef = best_trial[2]
        best_err = best_trial[0]

        if not improved:
            break

    return best_coef, best_err


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--method", choices=["lp", "remez", "both"], default="both")
    parser.add_argument("--degree", type=int, default=2)
    parser.add_argument("--n", type=int, default=41)
    parser.add_argument("--seed", type=int, default=-1)
    parser.add_argument("--tnoise", type=str, default="z2")
    parser.add_argument("--knoise", type=float, default=0.8)
    args = parser.parse_args()

    x, y, y_true = make_data(args.n, args.tnoise, args.knoise, seed=args.seed)

    coef_l2 = least_squares_fit(x, y, args.degree)

    results = {"least squares": coef_l2}

    if args.method in ["lp", "both"]:
        coef_lp, err_lp = minimax_lp_fit(x, y, args.degree)
        results["minimax LP"] = coef_lp
        print("[LP minimax]")
        print("coef =", coef_lp)
        print("max abs error =", err_lp)

    if args.method in ["remez", "both"]:
        coef_remez, err_remez = minimax_remez_fit(x, y, args.degree)
        results["minimax Remez-like"] = coef_remez
        print("[Remez-like minimax]")
        print("coef =", coef_remez)
        print("max abs error =", err_remez)

    print("[least squares]")
    print("coef =", coef_l2)
    print("max abs error =", np.max(np.abs(predict(x, coef_l2) - y)))

    xx = np.linspace(0, 1, 400)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(9, 7), sharex=True)

    ax1.scatter(x, y, c="black", s=25, label="data")
    ax1.plot(xx, np.exp(2.2 * xx) - 1.0, "k:", label="true function")

    for name, coef in results.items():
        ax1.plot(xx, predict(xx, coef), label=name)

    ax1.set_ylabel("y")
    ax1.set_title(f"Minimax approximation vs least squares, degree={args.degree}")
    ax1.grid(True, alpha=0.3)
    ax1.legend()

    for name, coef in results.items():
        ax2.plot(x, predict(x, coef) - y, marker="o", label=name)

    ax2.axhline(0, color="gray", lw=1)
    ax2.set_xlabel("x")
    ax2.set_ylabel("residual")
    ax2.grid(True, alpha=0.3)
    ax2.legend()

    fig.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()