import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import platform

# Configure fonts for Chinese support
system = platform.system()
if system == 'Windows':
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial']
elif system == 'Darwin':
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'PingFang SC']
else:
    plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def rosenbrock(x, y, a=1, b=100):
    return (a - x)**2 + b * (y - x**2)**2

def grad_rosenbrock(x, y, a=1, b=100):
    dx = -2 * (a - x) - 4 * b * x * (y - x**2)
    dy = 2 * b * (y - x**2)
    return np.array([dx, dy])

# Optimizer Implementations

class SGD:
    def __init__(self, lr=0.001, momentum=0.9):
        self.lr = lr
        self.momentum = momentum
        self.v = np.zeros(2)

    def step(self, params, grads):
        self.v = self.momentum * self.v - self.lr * grads
        return params + self.v

class AdamW:
    def __init__(self, lr=0.01, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.01):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.wd = weight_decay
        self.m = np.zeros(2)
        self.v = np.zeros(2)
        self.t = 0

    def step(self, params, grads):
        self.t += 1
        # Weight decay
        params = params - self.lr * self.wd * params

        # Adam
        self.m = self.beta1 * self.m + (1 - self.beta1) * grads
        self.v = self.beta2 * self.v + (1 - self.beta2) * grads**2

        m_hat = self.m / (1 - self.beta1**self.t)
        v_hat = self.v / (1 - self.beta2**self.t)

        return params - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

class Muon:
    def __init__(self, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
        self.lr = lr
        self.momentum = momentum
        self.nesterov = nesterov
        self.ns_steps = ns_steps
        self.m = np.zeros(2)

    def newton_schulz(self, M):
        # M is shape (1, 2) or (2,) treated as (1, 2)
        # For a vector, this is effectively normalization, but let's simulate the matrix steps
        # Shape handling
        if M.ndim == 1:
            M = M.reshape(1, -1)

        # Normalization (Spectral norm approx via Frobenius for small matrix)
        # X = M / (||M|| + eps)
        norm = np.linalg.norm(M) + 1e-8
        X = M / norm

        # Newton-Schulz Iteration
        # X_{k+1} = 0.5 * X_k * (3I - X_k^T * X_k)
        # Note: For non-square matrices A (m x n), we use A * (3I - A^T A) or (3I - A A^T) * A depending on dimensions
        # Here 1x2, let's use the standard form for whitening: X * (3I - X.T @ X)
        # Wait, if X is 1x2, X.T @ X is 2x2. X @ X.T is 1x1.
        # The iteration converges to a matrix with singular values all 1.

        for _ in range(self.ns_steps):
            # For 1x2:
            # A = X (1x2)
            # B = 3I - X^T @ X (2x2)
            # This direction is computationally heavier for 1x2.
            # Alternative: X_{k+1} = 1.5 * X_k - 0.5 * X_k @ (X_k.T @ X_k)
            # Actually for vector v, v^T v is a scalar if column vector, but here row.

            XTX = X.T @ X # (2, 2)
            eye = np.eye(2)
            X = 0.5 * X @ (3 * eye - XTX)

        return X.flatten()

    def step(self, params, grads):
        self.m = self.momentum * self.m + grads

        if self.nesterov:
            v = self.momentum * self.m + grads
        else:
            v = self.m

        # Orthogonalize v using Newton-Schulz
        update = self.newton_schulz(v)

        # Scale by root-mean-square of params (often used in Muon implementations for consistency)
        # or just standard LR. Original Muon paper/impl often scales by max(1, rms(params))
        # Let's stick to simple LR for visualization

        return params - self.lr * update

def visualize():
    # Setup plot
    fig, ax = plt.subplots(figsize=(10, 8))

    # Generate meshgrid for contour
    x = np.linspace(-2, 2, 400)
    y = np.linspace(-1, 3, 400)
    X, Y = np.meshgrid(x, y)
    Z = rosenbrock(X, Y)

    # Plot contours (log scale for better visibility)
    ax.contour(X, Y, Z, levels=np.logspace(-1, 3, 30), cmap='gray_r', alpha=0.4)
    ax.plot(1, 1, 'r*', markersize=15, label='Global Min (1, 1)')

    # Starting point
    start_pos = np.array([-1.5, 2.5])

    # Run Optimizers
    optimizers = [
        ('SGD (Momentum)', SGD(lr=0.0005, momentum=0.9), 'blue'),
        ('AdamW', AdamW(lr=0.05, weight_decay=0.01), 'green'),
        ('Muon', Muon(lr=0.05, momentum=0.95), 'purple')
    ]

    steps = 200

    for name, opt, color in optimizers:
        path = [start_pos]
        params = start_pos.copy()

        for _ in range(steps):
            grads = grad_rosenbrock(params[0], params[1])
            params = opt.step(params, grads)
            path.append(params)

        path = np.array(path)
        ax.plot(path[:, 0], path[:, 1], '.-', color=color, label=name, alpha=0.8, linewidth=1.5, markersize=3)
        # Mark end
        ax.plot(path[-1, 0], path[-1, 1], 'o', color=color)

    ax.set_title("优化器寻优路径对比: SGD vs AdamW vs Muon\n(Loss: Rosenbrock Function)", fontsize=14)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.legend()
    ax.grid(True, linestyle='--', alpha=0.3)

    plt.tight_layout()
    plt.savefig('optimizer_comparison.png', dpi=150)
    print("Optimization comparison plot saved to optimizer_comparison.png")

if __name__ == "__main__":
    visualize()