"""
Optimizers — comprehensive topic content with interactive graphs.
"""

import json
import numpy as np


def _optimizer_paths():
    """Show different optimizer paths on a 2D loss surface."""
    np.random.seed(42)
    # Simulate paths on a Rosenbrock-like surface
    t = np.linspace(0, 1, 50)

    sgd_x = 3.0 * (1 - t) ** 2 + 1 * t ** 2 + np.random.normal(0, 0.15, 50) * (1 - t)
    sgd_y = 3.0 * (1 - t) ** 1.5 + 1 * t ** 1.5 + np.random.normal(0, 0.15, 50) * (1 - t)

    momentum_x = 3.0 * (1 - t) ** 1.3 + 1 * t ** 1.3 + np.random.normal(0, 0.08, 50) * (1 - t)
    momentum_y = 3.0 * (1 - t) ** 1.2 + 1 * t ** 1.2 + np.random.normal(0, 0.08, 50) * (1 - t)

    adam_x = 3.0 * (1 - t) ** 1.1 + 1 * t ** 1.1 + np.random.normal(0, 0.03, 50) * (1 - t)
    adam_y = 3.0 * (1 - t) ** 1.0 + 1 * t ** 1.0 + np.random.normal(0, 0.03, 50) * (1 - t)

    data = [
        {"x": [1], "y": [1], "mode": "markers", "type": "scatter", "name": "Optimum",
         "marker": {"size": 16, "color": "#fbbf24", "symbol": "star", "line": {"width": 2, "color": "#fff"}}},
        {"x": sgd_x.tolist(), "y": sgd_y.tolist(), "mode": "lines+markers", "type": "scatter", "name": "SGD",
         "line": {"color": "#FF6584", "width": 2}, "marker": {"size": 3}},
        {"x": momentum_x.tolist(), "y": momentum_y.tolist(), "mode": "lines+markers", "type": "scatter", "name": "SGD + Momentum",
         "line": {"color": "#34d399", "width": 2}, "marker": {"size": 3}},
        {"x": adam_x.tolist(), "y": adam_y.tolist(), "mode": "lines+markers", "type": "scatter", "name": "Adam",
         "line": {"color": "#6C63FF", "width": 2.5}, "marker": {"size": 3}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Optimizer Convergence Paths", "font": {"size": 18}},
        "xaxis": {"title": "Weight 1", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Weight 2", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def _lr_schedules():
    """Show learning rate schedules."""
    epochs = list(range(100))
    constant = [0.01] * 100
    step = [0.01 * (0.1 ** (e // 30)) for e in epochs]
    cosine = [0.01 * 0.5 * (1 + np.cos(np.pi * e / 100)) for e in epochs]
    warmup_cosine = [0.01 * min(1, e / 10) * 0.5 * (1 + np.cos(np.pi * max(0, e - 10) / 90)) for e in epochs]

    data = [
        {"x": epochs, "y": constant, "mode": "lines", "type": "scatter", "name": "Constant",
         "line": {"color": "#FF6584", "width": 2, "dash": "dash"}},
        {"x": epochs, "y": step, "mode": "lines", "type": "scatter", "name": "Step Decay",
         "line": {"color": "#fbbf24", "width": 2.5}},
        {"x": epochs, "y": cosine, "mode": "lines", "type": "scatter", "name": "Cosine Annealing",
         "line": {"color": "#6C63FF", "width": 2.5}},
        {"x": epochs, "y": warmup_cosine, "mode": "lines", "type": "scatter", "name": "Warmup + Cosine",
         "line": {"color": "#34d399", "width": 3}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Learning Rate Schedules", "font": {"size": 18}},
        "xaxis": {"title": "Epoch", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Learning Rate", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def _convergence_speed():
    """Compare convergence speed of different optimizers."""
    np.random.seed(42)
    epochs = list(range(1, 51))
    sgd = [2.0 * np.exp(-0.04 * e) + 0.3 + np.random.normal(0, 0.04) for e in epochs]
    momentum = [2.0 * np.exp(-0.07 * e) + 0.15 + np.random.normal(0, 0.03) for e in epochs]
    rmsprop = [2.0 * np.exp(-0.09 * e) + 0.1 + np.random.normal(0, 0.025) for e in epochs]
    adam = [2.0 * np.exp(-0.12 * e) + 0.08 + np.random.normal(0, 0.02) for e in epochs]

    data = [
        {"x": epochs, "y": sgd, "mode": "lines", "type": "scatter", "name": "SGD",
         "line": {"color": "#FF6584", "width": 2.5}},
        {"x": epochs, "y": momentum, "mode": "lines", "type": "scatter", "name": "SGD + Momentum",
         "line": {"color": "#fbbf24", "width": 2.5}},
        {"x": epochs, "y": rmsprop, "mode": "lines", "type": "scatter", "name": "RMSProp",
         "line": {"color": "#a78bfa", "width": 2.5}},
        {"x": epochs, "y": adam, "mode": "lines", "type": "scatter", "name": "Adam",
         "line": {"color": "#34d399", "width": 3}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Convergence Speed Comparison", "font": {"size": 18}},
        "xaxis": {"title": "Epoch", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Loss", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 50},
    }
    return json.dumps({"data": data, "layout": layout})


def get_content() -> dict:
    return {
        "title": "Optimizers",
        "icon": "🏃",
        "subtitle": "Algorithms that minimize the loss function",
        "sections": [
            {"id": "introduction", "heading": "What are Optimizers?", "type": "text",
             "content": ("Optimizers use <strong>gradients</strong> computed by backpropagation to <strong>update model weights</strong> "
                         "in the direction that reduces the loss function.\n\n"
                         "The simplest is SGD (Stochastic Gradient Descent), but modern optimizers like <strong>Adam</strong> "
                         "use momentum and adaptive learning rates for much faster convergence.")},
            {"id": "paths", "heading": "Optimizer Convergence Paths", "type": "graph",
             "description": "Different optimizers take different paths to the minimum. SGD oscillates, momentum smooths the path, and Adam takes the most direct route with less noise.",
             "graph_json": _optimizer_paths()},
            {"id": "math", "heading": "The Mathematics", "type": "math",
             "content": [
                 {"label": "SGD", "formula": "\\theta_{t+1} = \\theta_t - \\eta \\nabla L(\\theta_t)",
                  "explanation": "Simplest optimizer. Updates weights by subtracting gradient × learning rate. Noisy but with good generalization."},
                 {"label": "Momentum", "formula": "v_t = \\beta v_{t-1} + \\nabla L, \\quad \\theta_{t+1} = \\theta_t - \\eta v_t",
                  "explanation": "Accumulates past gradients like a rolling ball. Accelerates in consistent directions, dampens oscillations. β=0.9 typical."},
                 {"label": "RMSProp", "formula": "s_t = \\beta s_{t-1} + (1-\\beta)(\\nabla L)^2, \\quad \\theta_{t+1} = \\theta_t - \\frac{\\eta}{\\sqrt{s_t + \\epsilon}} \\nabla L",
                  "explanation": "Adapts learning rate per-parameter. Parameters with large gradients get smaller updates. Handles sparse gradients well."},
                 {"label": "Adam", "formula": "m_t = \\beta_1 m_{t-1} + (1-\\beta_1)g, \\quad v_t = \\beta_2 v_{t-1} + (1-\\beta_2)g^2",
                  "explanation": "Combines momentum (m) and RMSProp (v) with bias correction. Default: β₁=0.9, β₂=0.999, η=0.001."},
             ]},
            {"id": "convergence", "heading": "Convergence Speed", "type": "graph",
             "description": "Adam converges fastest in most cases. SGD is slowest but can find flatter minima (better generalization). RMSProp and Momentum are in between.",
             "graph_json": _convergence_speed()},
            {"id": "schedules", "heading": "Learning Rate Schedules", "type": "graph",
             "description": "Starting with a higher LR for exploration and reducing it for fine-tuning. Warmup + Cosine is the modern standard for training transformers and large models.",
             "graph_json": _lr_schedules()},
            {"id": "comparison", "heading": "When to Use What", "type": "list",
             "entries": [
                 {"title": "Adam (Default)", "detail": "Works great out of the box. Default for most deep learning. lr=0.001, betas=(0.9, 0.999)."},
                 {"title": "AdamW", "detail": "Adam with decoupled weight decay. Correct regularization. Default in transformers and modern training."},
                 {"title": "SGD + Momentum + Schedule", "detail": "Often gives BEST final accuracy (flatter minima). Standard for computer vision. Requires careful LR scheduling."},
                 {"title": "RMSProp", "detail": "Good for RNNs. Handles non-stationary objectives well."},
                 {"title": "LAMB / LARS", "detail": "For very large batch training (thousands of GPUs). Layer-wise adaptive learning rates."},
             ]},
            {"id": "implementation", "heading": "Python Implementation", "type": "code", "language": "python",
             "content": ("import torch\nimport torch.optim as optim\n\nmodel = ...  # your model\n\n"
                         "# Adam (default choice)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\n"
                         "# AdamW (better regularization)\noptimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n\n"
                         "# SGD + Momentum (for best accuracy)\noptimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)\n\n"
                         "# Learning rate scheduler\nscheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)\n\n"
                         "# Training loop\nfor epoch in range(100):\n"
                         "    for batch in dataloader:\n"
                         "        loss = compute_loss(model, batch)\n"
                         "        optimizer.zero_grad()\n"
                         "        loss.backward()\n"
                         "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n"
                         "        optimizer.step()\n"
                         "    scheduler.step()  # update LR after each epoch\n")},
        ],
    }
