#!/usr/bin/env python
# optimize2d.py
#
# Usage examples:
#   python optimize2d.py --mode newton  --x0 0.0  --y0 0.0
#   python optimize2d.py --mode bfgs    --x0 -1.0 --y0 -1.0
#   python optimize2d.py --mode cg      --x0 -1.0 --y0 -1.0
#   python optimize2d.py --mode sd      --x0 0.0  --y0 1.0
#   python optimize2d.py --mode simplex --x0 0.0  --y0 1.0 --simplex_step 1.0
#   python optimize2d.py --mode simplex --x0 0.0 --y0 1.0 --simplex_step 1.0 --simplex_plot_interval 3
#   python optimize2d.py --mode simplex --x0 0.0  --y0 1.0
#   python optimize2d.py --mode all     --x0 -1.0 --y0 -1.0 --save 1 --show 1

from __future__ import annotations

import argparse
import sys
import traceback
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from dataclasses import dataclass, field
import matplotlib.pyplot as plt
from matplotlib import cm

try:
    from scipy.optimize import minimize
except Exception:
    minimize = None


# ============================================================
# Target function
# F(x,y) = -3 - 10x - 30x^2 + 1.5x^3 + 3x^4
#          + 30y - 30y^2 + 3y^4 + 3xy^2
# ============================================================

def func(v):
    x, y = float(v[0]), float(v[1])
    return (
        -3.0
        - 10.0 * x
        - 30.0 * x**2
        + 1.5 * x**3
        + 3.0 * x**4
        + 30.0 * y
        - 30.0 * y**2
        + 3.0 * y**4
        + 3.0 * x * y**2
    )


def grad(v):
    x, y = float(v[0]), float(v[1])
    dfdx = -10.0 - 60.0 * x + 4.5 * x**2 + 12.0 * x**3 + 3.0 * y**2
    dfdy = 30.0 - 60.0 * y + 12.0 * y**3 + 6.0 * x * y
    return np.array([dfdx, dfdy], dtype=float)


def hess(v):
    x, y = float(v[0]), float(v[1])
    return np.array(
        [
            [-60.0 + 9.0 * x + 36.0 * x**2, 6.0 * y],
            [6.0 * y, -60.0 + 36.0 * y**2 + 6.0 * x],
        ],
        dtype=float,
    )


# ============================================================
# Result container
# ============================================================

@dataclass
class OptResult:
    mode: str
    path: np.ndarray
    x: np.ndarray
    f: float
    grad_norm: float
    nit: int
    success: bool
    message: str
    simplexes: list = field(default_factory=list)


# ============================================================
# Optimizers
# ============================================================

def limit_step(dx, max_step):
    dx = np.asarray(dx, dtype=float)
    nrm = np.linalg.norm(dx)
    if max_step > 0.0 and nrm > max_step:
        dx = dx * (max_step / nrm)
    return dx


def run_newton(x0, args):
    """
    Newton-Raphson method for solving grad F = 0.

    注意:
      これは「最小化専用 Newton」ではなく、停留点を探す Newton 法です。
      そのため、初期値によって maximum / local minimum / global minimum に収束します。
    """
    x = np.asarray(x0, dtype=float)
    path = [x.copy()]
    success = False
    message = "maximum iteration reached"

    for it in range(1, args.nmaxiter + 1):
        g = grad(x)
        H = hess(x)

        if np.linalg.norm(g) < args.tol:
            success = True
            message = "gradient norm converged"
            break

        try:
            dx = -np.linalg.solve(H, g)
        except np.linalg.LinAlgError:
            dx = -np.linalg.pinv(H) @ g

        dx = limit_step(dx, args.max_step)
        x = x + args.damping * dx
        path.append(x.copy())

        if args.print_interval > 0 and it % args.print_interval == 0:
            print(
                f"[newton] iter={it:4d} "
                f"x={x[0]: .8f} y={x[1]: .8f} "
                f"F={func(x): .8e} |grad|={np.linalg.norm(grad(x)): .3e}"
            )

        if np.linalg.norm(dx) < args.tol:
            success = True
            message = "step norm converged"
            break

    return OptResult(
        mode="newton",
        path=np.asarray(path),
        x=x,
        f=float(func(x)),
        grad_norm=float(np.linalg.norm(grad(x))),
        nit=len(path) - 1,
        success=success,
        message=message,
    )


def run_sd(x0, args):
    """
    Steepest descent with fixed alpha.

    linesearch は使わない。
    講義用に、軌跡が見やすい単純な実装にしている。
    """
    x = np.asarray(x0, dtype=float)
    path = [x.copy()]
    success = False
    message = "maximum iteration reached"

    for it in range(1, args.nmaxiter + 1):
        g = grad(x)

        if np.linalg.norm(g) < args.tol:
            success = True
            message = "gradient norm converged"
            break

        dx = -args.alpha * g
        dx = limit_step(dx, args.max_step)
        x = x + dx
        path.append(x.copy())

        if args.print_interval > 0 and it % args.print_interval == 0:
            print(
                f"[sd] iter={it:4d} "
                f"x={x[0]: .8f} y={x[1]: .8f} "
                f"F={func(x): .8e} |grad|={np.linalg.norm(grad(x)): .3e}"
            )

        if np.linalg.norm(dx) < args.tol:
            success = True
            message = "step norm converged"
            break

    return OptResult(
        mode="sd",
        path=np.asarray(path),
        x=x,
        f=float(func(x)),
        grad_norm=float(np.linalg.norm(grad(x))),
        nit=len(path) - 1,
        success=success,
        message=message,
    )


def run_simplex(x0, args):
    """
    Nelder-Mead simplex method for 2D minimization.

    SciPy の callback では simplex（三角形）の各頂点が取れないので、
    講義用可視化のために最小限の Nelder-Mead を自前実装する。
    """
    x0 = np.asarray(x0, dtype=float)
    n = 2

    step = float(args.simplex_step)
    if step <= 0:
        raise ValueError("--simplex_step must be positive")

    simplex = np.vstack([
        x0,
        x0 + np.array([step, 0.0]),
        x0 + np.array([0.0, step]),
    ])

    alpha = 1.0   # reflection
    gamma = 2.0   # expansion
    rho   = 0.5   # contraction
    sigma = 0.5   # shrink

    path = [x0.copy()]
    simplexes = [simplex.copy()]

    success = False
    message = "maximum iteration reached"

    for it in range(1, args.nmaxiter + 1):
        fvals = np.array([func(v) for v in simplex])
        order = np.argsort(fvals)
        simplex = simplex[order]
        fvals = fvals[order]

        best = simplex[0]
        worst = simplex[-1]

        simplex_size = np.max(np.linalg.norm(simplex - best, axis=1))
        f_spread = abs(fvals[-1] - fvals[0])

        if simplex_size < args.tol and f_spread < args.tol:
            success = True
            message = "simplex size and function spread converged"
            break

        centroid = np.mean(simplex[:-1], axis=0)

        # reflection
        xr = centroid + alpha * (centroid - worst)
        fr = func(xr)

        if fvals[0] <= fr < fvals[-2]:
            simplex[-1] = xr

        elif fr < fvals[0]:
            # expansion
            xe = centroid + gamma * (xr - centroid)
            fe = func(xe)
            simplex[-1] = xe if fe < fr else xr

        else:
            # contraction
            if fr < fvals[-1]:
                # outside contraction
                xc = centroid + rho * (xr - centroid)
                fc = func(xc)
                if fc <= fr:
                    simplex[-1] = xc
                else:
                    for i in range(1, n + 1):
                        simplex[i] = best + sigma * (simplex[i] - best)
            else:
                # inside contraction
                xc = centroid + rho * (worst - centroid)
                fc = func(xc)
                if fc < fvals[-1]:
                    simplex[-1] = xc
                else:
                    for i in range(1, n + 1):
                        simplex[i] = best + sigma * (simplex[i] - best)

        fvals = np.array([func(v) for v in simplex])
        order = np.argsort(fvals)
        simplex = simplex[order]
        fvals = fvals[order]

        path.append(simplex[0].copy())
        simplexes.append(simplex.copy())

        if args.print_interval > 0 and it % args.print_interval == 0:
            print(
                f"[simplex] iter={it:4d} "
                f"x={simplex[0,0]: .8f} y={simplex[0,1]: .8f} "
                f"F={fvals[0]: .8e} "
                f"size={simplex_size:.3e}"
            )

    x = simplex[0].copy()

    return OptResult(
        mode="simplex",
        path=np.asarray(path),
        x=x,
        f=float(func(x)),
        grad_norm=float(np.linalg.norm(grad(x))),
        nit=len(path) - 1,
        success=success,
        message=message,
        simplexes=simplexes,
    )
    
def run_scipy(mode, x0, args):
    if minimize is None:
        raise ImportError("scipy is required for bfgs, cg, and simplex modes.")

    scipy_method = {
        "bfgs": "BFGS",
        "cg": "CG",
#        "simplex": "Nelder-Mead",
    }[mode]

    x0 = np.asarray(x0, dtype=float)
    path = [x0.copy()]

    def callback(xk):
        xk = np.asarray(xk, dtype=float)
        path.append(xk.copy())

        it = len(path) - 1
        if args.print_interval > 0 and it % args.print_interval == 0:
            print(
                f"[{mode}] iter={it:4d} "
                f"x={xk[0]: .8f} y={xk[1]: .8f} "
                f"F={func(xk): .8e} |grad|={np.linalg.norm(grad(xk)): .3e}"
            )

    options = {"maxiter": args.nmaxiter, "disp": False}

    if mode in ["bfgs", "cg"]:
        options["gtol"] = args.tol
        res = minimize(
            func,
            x0,
            jac=grad,
            method=scipy_method,
            callback=callback,
            tol=args.tol,
            options=options,
        )
    else:
        options["xatol"] = args.tol
        options["fatol"] = args.tol
        res = minimize(
            func,
            x0,
            method=scipy_method,
            callback=callback,
            tol=args.tol,
            options=options,
        )

    if np.linalg.norm(path[-1] - res.x) > 1.0e-14:
        path.append(np.asarray(res.x, dtype=float).copy())

    return OptResult(
        mode=mode,
        path=np.asarray(path),
        x=np.asarray(res.x, dtype=float),
        f=float(res.fun),
        grad_norm=float(np.linalg.norm(grad(res.x))),
        nit=int(getattr(res, "nit", len(path) - 1)),
        success=bool(res.success),
        message=str(res.message),
    )


def run_optimizer(mode, x0, args):
    if mode == "newton":
        return run_newton(x0, args)
    if mode == "sd":
        return run_sd(x0, args)
    if mode == "simplex":
        return run_simplex(x0, args)
    if mode in ["bfgs", "cg"]:
        return run_scipy(mode, x0, args)
    raise ValueError(f"unknown mode: {mode}")


# ============================================================
# Plot utilities
# ============================================================

def make_grid(args):
    xs = np.linspace(args.xmin, args.xmax, args.ngrid)
    ys = np.linspace(args.ymin, args.ymax, args.ngrid)
    X, Y = np.meshgrid(xs, ys)

    Z = (
        -3.0
        - 10.0 * X
        - 30.0 * X**2
        + 1.5 * X**3
        + 3.0 * X**4
        + 30.0 * Y
        - 30.0 * Y**2
        + 3.0 * Y**4
        + 3.0 * X * Y**2
    )
    return X, Y, Z


def classify_stationary_point(x, X, Y, Z):
    eig = np.linalg.eigvalsh(hess(x))

    if np.all(eig > 1.0e-6):
        idx = np.unravel_index(np.argmin(Z), Z.shape)
        xg = np.array([X[idx], Y[idx]], dtype=float)
        if np.linalg.norm(x - xg) < 0.35:
            return "global minimum"
        return "local minimum"

    if np.all(eig < -1.0e-6):
        return "maximum"

    return "saddle"


def draw_surface(ax, X, Y, Z):
    ax.plot_wireframe(
        X,
        Y,
        Z,
        rstride=2,
        cstride=2,
        linewidth=0.6,
        color="tab:blue",
        alpha=0.9,
    )
    ax.set_title("surface")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("F")

def draw_simplex_triangles(ax, result, args):
    if result.mode != "simplex":
        return
    if not result.simplexes:
        return

    simplexes = result.simplexes
    n = len(simplexes)

    interval = max(1, int(args.simplex_plot_interval))

    for i, tri in enumerate(simplexes):
        if i % interval != 0 and i != n - 1:
            continue

        tri = np.asarray(tri, dtype=float)
        tri_closed = np.vstack([tri, tri[0]])

        t = i / max(n - 1, 1)

        ax.plot(
            tri_closed[:, 0],
            tri_closed[:, 1],
            "-",
            color=plt.cm.plasma(t),
            lw=1.0 + 1.0 * t,
            alpha=0.25 + 0.55 * t,
            zorder=3,
        )

    # 初期 simplex を強調
    tri0 = np.asarray(simplexes[0], dtype=float)
    tri0 = np.vstack([tri0, tri0[0]])
    ax.plot(
        tri0[:, 0],
        tri0[:, 1],
        "--",
        color="white",
        lw=1.8,
        alpha=0.9,
        zorder=4,
        label="initial simplex",
    )

    # 最終 simplex を強調
    trif = np.asarray(simplexes[-1], dtype=float)
    trif = np.vstack([trif, trif[0]])
    ax.plot(
        trif[:, 0],
        trif[:, 1],
        "-",
        color="cyan",
        lw=2.2,
        alpha=0.95,
        zorder=5,
        label="final simplex",
    )
    
def draw_contour(ax, result, X, Y, Z, args):
    label = classify_stationary_point(result.x, X, Y, Z)

    ax.contourf(X, Y, Z, levels=args.nlevels, cmap=args.cmap)
    ax.contour(X, Y, Z, levels=20, colors="white", linewidths=0.4, alpha=0.25)

    # Simplex 法の場合だけ、三角形の変化を描く
    draw_simplex_triangles(ax, result, args)

    p = result.path
    ax.plot(p[:, 0], p[:, 1], "-o", color="blue", lw=1.4, ms=3, zorder=6)
    ax.plot(p[0, 0], p[0, 1], "o", color="white", mec="black", ms=7, label="start", zorder=7)
    ax.plot(p[-1, 0], p[-1, 1], "o", color="cyan", mec="black", ms=7, label="end", zorder=7)

    ax.text(
        0.05,
        0.08,
        label,
        transform=ax.transAxes,
        color="white",
        fontsize=13,
        weight="bold",
    )

    ax.set_xlim(args.xmin, args.xmax)
    ax.set_ylim(args.ymin, args.ymax)
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_title(
        f"{result.mode}: {label}\n"
        f"end=({result.x[0]:.4g}, {result.x[1]:.4g}), F={result.f:.5g}"
    )

def plot_single(result, args):
    X, Y, Z = make_grid(args)

    fig = plt.figure(figsize=(12, 5))
    ax1 = fig.add_subplot(1, 2, 1, projection="3d")
    ax2 = fig.add_subplot(1, 2, 2)

    draw_surface(ax1, X, Y, Z)
    draw_contour(ax2, result, X, Y, Z, args)

    fig.suptitle(
        r"$F(x,y)=-3-10x-30x^2+1.5x^3+3x^4+30y-30y^2+3y^4+3xy^2$",
        fontsize=14,
    )
    fig.tight_layout()
    return fig


def plot_comparison(results, args):
    X, Y, Z = make_grid(args)

    fig = plt.figure(figsize=(15, 9))

    ax0 = fig.add_subplot(2, 3, 1, projection="3d")
    draw_surface(ax0, X, Y, Z)

    for i, result in enumerate(results, start=2):
        ax = fig.add_subplot(2, 3, i)
        draw_contour(ax, result, X, Y, Z, args)

    fig.suptitle(
        "Comparison: Newton, BFGS, CG, SD, Simplex",
        fontsize=20,
        color="blue",
        weight="bold",
    )
    fig.tight_layout()
    return fig


def save_figure(fig, filename):
    filename = Path(filename)
    filename.parent.mkdir(parents=True, exist_ok=True)

    print(f"[save] writing {filename}")
    fig.savefig(filename, dpi=160, bbox_inches="tight")
    print(f"[save] done {filename}")


# ============================================================
# CLI
# ============================================================

def build_parser():
    parser = argparse.ArgumentParser(
        description="2D nonlinear optimization demo with mode=newton/bfgs/cg/sd/simplex/all"
    )

    parser.add_argument(
        "--mode",
        required=True,
        choices=["newton", "bfgs", "cg", "sd", "simplex", "all"],
        help="optimization mode",
    )

    parser.add_argument("--infile", default="none", help="unused dummy argument")
    parser.add_argument("--outfile", default="", help="output prefix")

    parser.add_argument("--x0", type=float, default=0.0, help="initial x")
    parser.add_argument("--y0", type=float, default=0.0, help="initial y")

    parser.add_argument("--nmaxiter", type=int, default=200, help="maximum iterations")
    parser.add_argument("--tol", type=float, default=1.0e-5, help="convergence tolerance")

    parser.add_argument("--alpha", type=float, default=0.01, help="fixed step size for SD")
    parser.add_argument("--damping", type=float, default=1.0, help="damping factor for Newton")
    parser.add_argument("--max_step", type=float, default=2.0, help="maximum step length; <=0 disables clipping")

    parser.add_argument(
        "--simplex_step",
        type=float,
        default=2.0,
        help="initial simplex size for mode=simplex",
    )

    parser.add_argument(
        "--simplex_plot_interval",
        type=int,
        default=1,
        help="plot every N simplex triangles",
    )
    parser.add_argument("--xmin", type=float, default=-4.0)
    parser.add_argument("--xmax", type=float, default=4.0)
    parser.add_argument("--ymin", type=float, default=-4.0)
    parser.add_argument("--ymax", type=float, default=4.0)
    parser.add_argument("--ngrid", type=int, default=101)
    parser.add_argument("--nlevels", type=int, default=50)
    parser.add_argument("--cmap", default="turbo")

    parser.add_argument("--save", type=int, choices=[0, 1], default=1)
    parser.add_argument("--show", type=int, choices=[0, 1], default=1)

    parser.add_argument("--print_interval", type=int, default=1)

    return parser


def make_outfile(args):
    if args.outfile:
        prefix = args.outfile
    else:
        prefix = f"optimize2d_{args.mode}_x{args.x0:g}_y{args.y0:g}"

    prefix = (
        prefix.replace("-", "m")
        .replace("+", "p")
        .replace(".", "p")
        .replace(" ", "_")
    )

    return f"{prefix}.png"


def print_summary(result):
    print("")
    print(f"[result] mode      = {result.mode}")
    print(f"[result] success   = {result.success}")
    print(f"[result] message   = {result.message}")
    print(f"[result] nit       = {result.nit}")
    print(f"[result] x         = {result.x[0]:.12g}")
    print(f"[result] y         = {result.x[1]:.12g}")
    print(f"[result] F         = {result.f:.12g}")
    print(f"[result] |grad|    = {result.grad_norm:.12g}")


def main():
    args = build_parser().parse_args()

    x0 = np.array([args.x0, args.y0], dtype=float)

    if args.mode == "all":
        modes = ["newton", "bfgs", "cg", "sd", "simplex"]
        results = []

        for mode in modes:
            print("")
            print(f"[run] mode={mode}")
            result = run_optimizer(mode, x0, args)
            print_summary(result)
            results.append(result)

        fig = plot_comparison(results, args)

    else:
        print(f"[run] mode={args.mode}")
        result = run_optimizer(args.mode, x0, args)
        print_summary(result)
        fig = plot_single(result, args)

    if args.save == 1:
        save_figure(fig, make_outfile(args))

    if args.show == 1:
        plt.show()
    else:
        plt.close(fig)


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        sys.exit(1)