from __future__ import annotations

import logging
import os

# Import gradient descent utilities
import sys

import numpy as np
import sympy as sp

import QuizGenerator.contentast as ca
from QuizGenerator.mixins import BodyTemplatesMixin, TableQuestionMixin
from QuizGenerator.question import Question, QuestionRegistry

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'gradient_descent'))
from misc import format_vector, generate_function

log = logging.getLogger(__name__)


@QuestionRegistry.register()
class ParameterCountingQuestion(Question):
  """
  Question asking students to count parameters in a neural network.

  Given a dense network architecture, students calculate:
  - Total number of weights
  - Total number of biases
  - Total trainable parameters
  """

  def __init__(self, *args, **kwargs):
    kwargs["topic"] = kwargs.get("topic", Question.Topic.ML_OPTIMIZATION)
    super().__init__(*args, **kwargs)

    self.num_layers = kwargs.get("num_layers", None)
    self.include_biases = kwargs.get("include_biases", True)

  @classmethod
  def _build_context(cls, *, rng_seed=None, **kwargs):
    context = super()._build_context(rng_seed=rng_seed, **kwargs)
    context.num_layers = kwargs.get("num_layers")
    context.include_biases = kwargs.get("include_biases", True)

    # Generate random architecture
    if context.num_layers is None:
      context.num_layers = context.rng.choice([3, 4])

    # Generate layer sizes
    # Input layer: common sizes for typical problems
    input_sizes = [28*28, 32*32, 784, 1024, 64, 128]
    context.layer_sizes = [context.rng.choice(input_sizes)]

    # Hidden layers: reasonable sizes
    for _ in range(context.num_layers - 2):
      hidden_size = context.rng.choice([32, 64, 128, 256, 512])
      context.layer_sizes.append(hidden_size)

    # Output layer: typical classification sizes
    output_size = context.rng.choice([2, 10, 100, 1000])
    context.layer_sizes.append(output_size)

    # Calculate correct answers
    context.total_weights = 0
    context.total_biases = 0
    context.weights_per_layer = []
    context.biases_per_layer = []

    for i in range(len(context.layer_sizes) - 1):
      weights = context.layer_sizes[i] * context.layer_sizes[i+1]
      biases = context.layer_sizes[i+1] if context.include_biases else 0

      context.weights_per_layer.append(weights)
      context.biases_per_layer.append(biases)

      context.total_weights += weights
      context.total_biases += biases

    context.total_params = context.total_weights + context.total_biases
    return context

  @classmethod
  def _build_body(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question body and collect answers."""
    body = ca.Section()
    answers = []

    # Question description
    body.add_element(ca.Paragraph([
      "Consider a fully-connected (dense) neural network with the following architecture:"
    ]))

    # Display architecture
    arch_parts = []
    for i, size in enumerate(context.layer_sizes):
      if i > 0:
        arch_parts.append(" → ")
      arch_parts.append(str(size))

    body.add_element(ca.Paragraph([
      "Architecture: " + "".join(arch_parts)
    ]))

    if context.include_biases:
      body.add_element(ca.Paragraph([
        "Each layer includes bias terms."
      ]))

    # Questions
    # Answer table
    table_data = []
    table_data.append(["Parameter Type", "Count"])

    total_weights_answer = ca.AnswerTypes.Int(context.total_weights, label="Total weights")
    total_biases_answer = None
    total_params_answer = ca.AnswerTypes.Int(context.total_params, label="Total trainable parameters")

    answers.append(total_weights_answer)
    table_data.append([
      "Total weights (connections between layers)",
      total_weights_answer
    ])

    if context.include_biases:
      total_biases_answer = ca.AnswerTypes.Int(context.total_biases, label="Total biases")
      answers.append(total_biases_answer)
      table_data.append([
        "Total biases",
        total_biases_answer
      ])

    answers.append(total_params_answer)
    table_data.append([
      "Total trainable parameters",
      total_params_answer
    ])

    body.add_element(ca.Table(data=table_data))

    return body, answers

  @classmethod
  def _build_explanation(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question explanation."""
    explanation = ca.Section()

    explanation.add_element(ca.Paragraph([
      "To count parameters in a dense neural network, we calculate weights and biases for each layer."
    ]))

    explanation.add_element(ca.Paragraph([
      ca.Text("Weights calculation:", emphasis=True)
    ]))

    for i in range(len(context.layer_sizes) - 1):
      input_size = context.layer_sizes[i]
      output_size = context.layer_sizes[i+1]
      weights = context.weights_per_layer[i]

      explanation.add_element(ca.Paragraph([
        f"Layer {i+1} → {i+2}: ",
        ca.Equation(f"{input_size} \\times {output_size} = {weights:,}", inline=True),
        " weights"
      ]))

    explanation.add_element(ca.Paragraph([
      "Total weights: ",
      ca.Equation(
        f"{' + '.join([f'{w:,}' for w in context.weights_per_layer])} = {context.total_weights:,}",
        inline=True
      )
    ]))

    if context.include_biases:
      explanation.add_element(ca.Paragraph([
        ca.Text("Biases calculation:", emphasis=True)
      ]))

      for i in range(len(context.layer_sizes) - 1):
        output_size = context.layer_sizes[i+1]
        biases = context.biases_per_layer[i]

        explanation.add_element(ca.Paragraph([
          f"Layer {i+2}: {biases:,} biases (one per neuron)"
        ]))

      explanation.add_element(ca.Paragraph([
        "Total biases: ",
        ca.Equation(
          f"{' + '.join([f'{b:,}' for b in context.biases_per_layer])} = {context.total_biases:,}",
          inline=True
        )
      ]))

    explanation.add_element(ca.Paragraph([
      ca.Text("Total trainable parameters:", emphasis=True)
    ]))

    if context.include_biases:
      explanation.add_element(ca.Equation(
        f"\\text{{Total}} = {context.total_weights:,} + {context.total_biases:,} = {context.total_params:,}",
        inline=False
      ))
    else:
      explanation.add_element(ca.Equation(
        f"\\text{{Total}} = {context.total_weights:,}",
        inline=False
      ))

    return explanation, []


@QuestionRegistry.register()
class ActivationFunctionComputationQuestion(Question):
  """
  Question asking students to compute activation function outputs.

  Given a vector of inputs and an activation function, students calculate
  the output for each element (or entire vector for softmax).
  """

  ACTIVATION_RELU = "relu"
  ACTIVATION_SIGMOID = "sigmoid"
  ACTIVATION_TANH = "tanh"
  ACTIVATION_SOFTMAX = "softmax"

  def __init__(self, *args, **kwargs):
    kwargs["topic"] = kwargs.get("topic", Question.Topic.ML_OPTIMIZATION)
    super().__init__(*args, **kwargs)

    self.vector_size = kwargs.get("vector_size", None)
    self.activation = kwargs.get("activation", None)

  @classmethod
  def _build_context(cls, *, rng_seed=None, **kwargs):
    context = super()._build_context(rng_seed=rng_seed, **kwargs)
    context.vector_size = kwargs.get("vector_size")
    context.activation = kwargs.get("activation")

    # Generate random input vector
    if context.vector_size is None:
      context.vector_size = context.rng.choice([3, 4, 5])

    context.input_vector = [
      round(context.rng.uniform(-3, 3), 1)
      for _ in range(context.vector_size)
    ]

    # Select activation function
    if context.activation is None:
      activations = [
        cls.ACTIVATION_RELU,
        cls.ACTIVATION_SIGMOID,
        cls.ACTIVATION_TANH,
        cls.ACTIVATION_SOFTMAX,
      ]
      context.activation = context.rng.choice(activations)

    # For leaky ReLU, set alpha
    context.leaky_alpha = 0.01

    # Compute outputs
    context.output_vector = cls._compute_activation(context.activation, context.input_vector)
    return context

  @staticmethod
  def _compute_activation(activation, inputs):
    """Compute activation function output."""
    if activation == ActivationFunctionComputationQuestion.ACTIVATION_RELU:
      return [max(0, x) for x in inputs]

    elif activation == ActivationFunctionComputationQuestion.ACTIVATION_SIGMOID:
      return [1 / (1 + np.exp(-x)) for x in inputs]

    elif activation == ActivationFunctionComputationQuestion.ACTIVATION_TANH:
      return [np.tanh(x) for x in inputs]

    elif activation == ActivationFunctionComputationQuestion.ACTIVATION_SOFTMAX:
      # Subtract max for numerical stability
      exp_vals = [np.exp(x - max(inputs)) for x in inputs]
      sum_exp = sum(exp_vals)
      return [e / sum_exp for e in exp_vals]

    else:
      raise ValueError(f"Unknown activation: {activation}")

  @staticmethod
  def _get_activation_name(activation):
    """Get human-readable activation name."""
    names = {
      ActivationFunctionComputationQuestion.ACTIVATION_RELU: "ReLU",
      ActivationFunctionComputationQuestion.ACTIVATION_SIGMOID: "Sigmoid",
      ActivationFunctionComputationQuestion.ACTIVATION_TANH: "Tanh",
      ActivationFunctionComputationQuestion.ACTIVATION_SOFTMAX: "Softmax",
    }
    return names.get(activation, "Unknown")

  @staticmethod
  def _get_activation_formula(activation):
    """Get LaTeX formula for activation function."""
    if activation == ActivationFunctionComputationQuestion.ACTIVATION_RELU:
      return r"\text{ReLU}(x) = \max(0, x)"

    elif activation == ActivationFunctionComputationQuestion.ACTIVATION_SIGMOID:
      return r"\sigma(x) = \frac{1}{1 + e^{-x}}"

    elif activation == ActivationFunctionComputationQuestion.ACTIVATION_TANH:
      return r"\tanh(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}"

    elif activation == ActivationFunctionComputationQuestion.ACTIVATION_SOFTMAX:
      return r"\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}"

    return ""

  @classmethod
  def _build_body(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question body and collect answers."""
    body = ca.Section()
    answers = []

    # Question description
    body.add_element(ca.Paragraph([
      f"Given the input vector below, compute the output after applying the {cls._get_activation_name(context.activation)} activation function."
    ]))

    # Display formula
    body.add_element(ca.Paragraph([
      "Activation function: ",
      ca.Equation(cls._get_activation_formula(context.activation), inline=True)
    ]))

    # Input vector
    input_str = ", ".join([f"{x:.1f}" for x in context.input_vector])
    body.add_element(ca.Paragraph([
      "Input: ",
      ca.Equation(f"[{input_str}]", inline=True)
    ]))

    # Answer table
    if context.activation == cls.ACTIVATION_SOFTMAX:
      body.add_element(ca.Paragraph([
        "Compute the output vector:"
      ]))

      output_answer = ca.AnswerTypes.Vector(context.output_vector, label="Output vector")
      answers.append(output_answer)
      table_data = []
      table_data.append(["Output Vector"])
      table_data.append([output_answer])

      body.add_element(ca.Table(data=table_data))

    else:
      body.add_element(ca.Paragraph([
        "Compute the output for each element:"
      ]))

      table_data = []
      table_data.append(["Input", "Output"])

      for i, x in enumerate(context.input_vector):
        answer = ca.AnswerTypes.Float(
          float(context.output_vector[i]),
          label=f"Output for input {context.input_vector[i]:.1f}"
        )
        answers.append(answer)
        table_data.append([
          ca.Equation(f"{x:.1f}", inline=True),
          answer
        ])

      body.add_element(ca.Table(data=table_data))

    return body, answers

  @classmethod
  def _build_explanation(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question explanation."""
    explanation = ca.Section()

    explanation.add_element(ca.Paragraph([
      f"To compute the {cls._get_activation_name(context.activation)} activation, we apply the formula to each input."
    ]))

    if context.activation == cls.ACTIVATION_SOFTMAX:
      explanation.add_element(ca.Paragraph([
        ca.Text("Softmax computation:", emphasis=True)
      ]))

      # Show exponentials
      exp_strs = [f"e^{{{x:.1f}}}" for x in context.input_vector]
      explanation.add_element(ca.Paragraph([
        "First, compute exponentials: ",
        ca.Equation(", ".join(exp_strs), inline=True)
      ]))

      # Numerical values
      exp_vals = [np.exp(x) for x in context.input_vector]
      exp_vals_str = ", ".join([f"{e:.4f}" for e in exp_vals])
      explanation.add_element(ca.Paragraph([
        ca.Equation(f"\\approx [{exp_vals_str}]", inline=True)
      ]))

      # Sum
      sum_exp = sum(exp_vals)
      explanation.add_element(ca.Paragraph([
        "Sum: ",
        ca.Equation(f"{sum_exp:.4f}", inline=True)
      ]))

      # Final outputs
      explanation.add_element(ca.Paragraph([
        "Divide each by the sum:"
      ]))

      for i, (exp_val, output) in enumerate(zip(exp_vals, context.output_vector)):
        explanation.add_element(ca.Equation(
          f"\\text{{softmax}}({context.input_vector[i]:.1f}) = \\frac{{{exp_val:.4f}}}{{{sum_exp:.4f}}} = {output:.4f}",
          inline=False
        ))

    else:
      explanation.add_element(ca.Paragraph([
        ca.Text("Element-wise computation:", emphasis=True)
      ]))

      for i, (x, y) in enumerate(zip(context.input_vector, context.output_vector)):
        if context.activation == cls.ACTIVATION_RELU:
          explanation.add_element(ca.Equation(
            f"\\text{{ReLU}}({x:.1f}) = \\max(0, {x:.1f}) = {y:.4f}",
            inline=False
          ))

        elif context.activation == cls.ACTIVATION_SIGMOID:
          explanation.add_element(ca.Equation(
            f"\\sigma({x:.1f}) = \\frac{{1}}{{1 + e^{{-{x:.1f}}}}} = {y:.4f}",
            inline=False
          ))

        elif context.activation == cls.ACTIVATION_TANH:
          explanation.add_element(ca.Equation(
            f"\\tanh({x:.1f}) = {y:.4f}",
            inline=False
          ))

    return explanation, []


@QuestionRegistry.register()
class RegularizationCalculationQuestion(Question):
  """
  Question asking students to calculate loss with L2 regularization.

  Given a small network (2-4 weights), students calculate:
  - Forward pass
  - Base MSE loss
  - L2 regularization penalty
  - Total loss
  - Gradient with regularization for one weight
  """

  def __init__(self, *args, **kwargs):
    kwargs["topic"] = kwargs.get("topic", Question.Topic.ML_OPTIMIZATION)
    super().__init__(*args, **kwargs)

    self.num_weights = kwargs.get("num_weights", None)

  @classmethod
  def _build_context(cls, *, rng_seed=None, **kwargs):
    context = super()._build_context(rng_seed=rng_seed, **kwargs)
    context.num_weights = kwargs.get("num_weights")

    # Generate small network (2-4 weights for simplicity)
    if context.num_weights is None:
      context.num_weights = context.rng.choice([2, 3, 4])

    # Generate weights (small values)
    context.weights = [
      round(context.rng.uniform(-2, 2), 1)
      for _ in range(context.num_weights)
    ]

    # Generate input and target
    context.input_val = round(context.rng.uniform(-3, 3), 1)
    context.target = round(context.rng.uniform(-5, 5), 1)

    # Regularization coefficient
    context.lambda_reg = context.rng.choice([0.01, 0.05, 0.1, 0.5])

    # Forward pass (simple linear combination for simplicity)
    # prediction = sum(w_i * input^i) for i in 0..n
    # This gives us a polynomial: w0 + w1*x + w2*x^2 + ...
    context.prediction = sum(
      w * (context.input_val ** i)
      for i, w in enumerate(context.weights)
    )

    # Calculate losses
    context.base_loss = 0.5 * (context.target - context.prediction) ** 2
    context.l2_penalty = (context.lambda_reg / 2) * sum(w**2 for w in context.weights)
    context.total_loss = context.base_loss + context.l2_penalty

    # Calculate gradient for first weight (w0, the bias term)
    # dL_base/dw0 = -(target - prediction) * dPrediction/dw0
    # dPrediction/dw0 = input^0 = 1
    # dL_reg/dw0 = lambda * w0
    # dL_total/dw0 = dL_base/dw0 + dL_reg/dw0

    context.grad_base_w0 = -(context.target - context.prediction) * 1  # derivative of w0*x^0
    context.grad_reg_w0 = context.lambda_reg * context.weights[0]
    context.grad_total_w0 = context.grad_base_w0 + context.grad_reg_w0
    return context

  @classmethod
  def _build_body(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question body and collect answers."""
    body = ca.Section()
    answers = []

    # Question description
    body.add_element(ca.Paragraph([
      "Consider a simple model with the following parameters:"
    ]))

    # Display weights
    weight_strs = [f"w_{i} = {w:.1f}" for i, w in enumerate(context.weights)]
    body.add_element(ca.Paragraph([
      "Weights: ",
      ca.Equation(", ".join(weight_strs), inline=True)
    ]))

    # Model equation
    terms = []
    for i, w in enumerate(context.weights):
      if i == 0:
        terms.append(f"w_0")
      elif i == 1:
        terms.append(f"w_1 x")
      else:
        terms.append(f"w_{i} x^{i}")

    model_eq = " + ".join(terms)
    body.add_element(ca.Paragraph([
      "Model: ",
      ca.Equation(f"\\hat{{y}} = {model_eq}", inline=True)
    ]))

    # Data point
    body.add_element(ca.Paragraph([
      "Data point: ",
      ca.Equation(f"x = {context.input_val:.1f}, y = {context.target:.1f}", inline=True)
    ]))

    # Regularization
    body.add_element(ca.Paragraph([
      "L2 regularization coefficient: ",
      ca.Equation(f"\\lambda = {context.lambda_reg}", inline=True)
    ]))

    body.add_element(ca.Paragraph([
      "Calculate the following:"
    ]))

    # Answer table
    table_data = []
    table_data.append(["Calculation", "Value"])

    prediction_answer = ca.AnswerTypes.Float(float(context.prediction), label="Prediction ŷ")
    base_loss_answer = ca.AnswerTypes.Float(float(context.base_loss), label="Base MSE loss")
    l2_penalty_answer = ca.AnswerTypes.Float(float(context.l2_penalty), label="L2 penalty")
    total_loss_answer = ca.AnswerTypes.Float(float(context.total_loss), label="Total loss")
    grad_total_w0_answer = ca.AnswerTypes.Float(float(context.grad_total_w0), label="Gradient ∂L/∂w₀")

    answers.append(prediction_answer)
    table_data.append([
      ca.Paragraph(["Prediction ", ca.Equation(r"\hat{y}", inline=True)]),
      prediction_answer
    ])

    answers.append(base_loss_answer)
    table_data.append([
      ca.Paragraph(["Base MSE loss: ", ca.Equation(r"L_{base} = (1/2)(y - \hat{y})^2", inline=True)]),
      base_loss_answer
    ])

    answers.append(l2_penalty_answer)
    table_data.append([
      ca.Paragraph(["L2 penalty: ", ca.Equation(r"L_{reg} = (\lambda/2)\sum w_i^2", inline=True)]),
      l2_penalty_answer
    ])

    answers.append(total_loss_answer)
    table_data.append([
      ca.Paragraph(["Total loss: ", ca.Equation(r"L_{total} = L_{base} + L_{reg}", inline=True)]),
      total_loss_answer
    ])

    answers.append(grad_total_w0_answer)
    table_data.append([
      ca.Paragraph(["Gradient: ", ca.Equation(r"\frac{\partial L_{total}}{\partial w_0}", inline=True)]),
      grad_total_w0_answer
    ])

    body.add_element(ca.Table(data=table_data))

    return body, answers

  @classmethod
  def _build_explanation(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question explanation."""
    explanation = ca.Section()

    explanation.add_element(ca.Paragraph([
      "L2 regularization adds a penalty term to the loss function to prevent overfitting by keeping weights small."
    ]))

    # Step 1: Forward pass
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 1: Compute prediction", emphasis=True)
    ]))

    terms = []
    for i, w in enumerate(context.weights):
      if i == 0:
        terms.append(f"{w:.1f}")
      else:
        x_term = f"{context.input_val:.1f}^{i}" if i > 1 else f"{context.input_val:.1f}"
        terms.append(f"{w:.1f} \\times {x_term}")

    explanation.add_element(ca.Equation(
      f"\\hat{{y}} = {' + '.join(terms)} = {context.prediction:.4f}",
      inline=False
    ))

    # Step 2: Base loss
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 2: Compute base MSE loss", emphasis=True)
    ]))

    explanation.add_element(ca.Equation(
      f"L_{{base}} = \\frac{{1}}{{2}}(y - \\hat{{y}})^2 = \\frac{{1}}{{2}}({context.target:.1f} - {context.prediction:.4f})^2 = {context.base_loss:.4f}",
      inline=False
    ))

    # Step 3: L2 penalty
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 3: Compute L2 penalty", emphasis=True)
    ]))

    weight_squares = [f"{w:.1f}^2" for w in context.weights]
    sum_squares = sum(w**2 for w in context.weights)

    explanation.add_element(ca.Equation(
      f"L_{{reg}} = \\frac{{\\lambda}}{{2}} \\sum w_i^2 = \\frac{{{context.lambda_reg}}}{{2}}({' + '.join(weight_squares)}) = \\frac{{{context.lambda_reg}}}{{2}} \\times {sum_squares:.4f} = {context.l2_penalty:.4f}",
      inline=False
    ))

    # Step 4: Total loss
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 4: Compute total loss", emphasis=True)
    ]))

    explanation.add_element(ca.Equation(
      f"L_{{total}} = L_{{base}} + L_{{reg}} = {context.base_loss:.4f} + {context.l2_penalty:.4f} = {context.total_loss:.4f}",
      inline=False
    ))

    # Step 5: Gradient with regularization
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 5: Compute gradient with regularization", emphasis=True)
    ]))

    explanation.add_element(ca.Paragraph([
      ca.Equation(r"w_0", inline=True),
      " (the bias term):"
    ]))

    explanation.add_element(ca.Equation(
      f"\\frac{{\\partial L_{{base}}}}{{\\partial w_0}} = -(y - \\hat{{y}}) \\times 1 = -({context.target:.1f} - {context.prediction:.4f}) = {context.grad_base_w0:.4f}",
      inline=False
    ))

    explanation.add_element(ca.Equation(
      f"\\frac{{\\partial L_{{reg}}}}{{\\partial w_0}} = \\lambda w_0 = {context.lambda_reg} \\times {context.weights[0]:.1f} = {context.grad_reg_w0:.4f}",
      inline=False
    ))

    explanation.add_element(ca.Equation(
      f"\\frac{{\\partial L_{{total}}}}{{\\partial w_0}} = {context.grad_base_w0:.4f} + {context.grad_reg_w0:.4f} = {context.grad_total_w0:.4f}",
      inline=False
    ))

    explanation.add_element(ca.Paragraph([
      "The regularization term adds ",
      ca.Equation(f"\\lambda w_0 = {context.grad_reg_w0:.4f}", inline=True),
      " to the gradient, pushing the weight toward zero."
    ]))

    return explanation, []


@QuestionRegistry.register()
class MomentumOptimizerQuestion(Question, TableQuestionMixin, BodyTemplatesMixin):
  """
  Question asking students to perform gradient descent with momentum.

  Given a function, current weights, gradients, learning rate, and momentum coefficient,
  students calculate:
  - Velocity update using momentum
  - Weight update using the new velocity
  - Comparison to vanilla SGD (optional)
  """

  def __init__(self, *args, **kwargs):
    kwargs["topic"] = kwargs.get("topic", Question.Topic.ML_OPTIMIZATION)
    super().__init__(*args, **kwargs)

    self.num_variables = kwargs.get("num_variables", 2)
    self.show_vanilla_sgd = kwargs.get("show_vanilla_sgd", True)

  @classmethod
  def _build_context(cls, *, rng_seed=None, **kwargs):
    context = super()._build_context(rng_seed=rng_seed, **kwargs)
    context.num_variables = kwargs.get("num_variables", 2)
    context.show_vanilla_sgd = kwargs.get("show_vanilla_sgd", True)

    # Generate well-conditioned quadratic function
    context.variables, context.function, context.gradient_function, context.equation = generate_function(
      context.rng,
      context.num_variables,
      max_degree=2,
      use_quadratic=True
    )

    # Generate current weights (small integers)
    context.current_weights = [
      context.rng.choice([-2, -1, 0, 1, 2])
      for _ in range(context.num_variables)
    ]

    # Calculate gradient at current position
    subs_map = dict(zip(context.variables, context.current_weights))
    g_syms = context.gradient_function.subs(subs_map)
    context.gradients = [float(val) for val in g_syms]

    # Generate previous velocity (for momentum)
    # Start with small or zero velocity
    context.prev_velocity = [
      round(context.rng.uniform(-0.5, 0.5), 2)
      for _ in range(context.num_variables)
    ]

    # Hyperparameters
    context.learning_rate = context.rng.choice([0.01, 0.05, 0.1])
    context.momentum_beta = context.rng.choice([0.8, 0.9])

    # Calculate momentum updates
    # v_new = beta * v_old + (1 - beta) * gradient
    context.new_velocity = [
      context.momentum_beta * v_old + (1 - context.momentum_beta) * grad
      for v_old, grad in zip(context.prev_velocity, context.gradients)
    ]

    # w_new = w_old - alpha * v_new
    context.new_weights = [
      w - context.learning_rate * v
      for w, v in zip(context.current_weights, context.new_velocity)
    ]

    # Calculate vanilla SGD for comparison
    if context.show_vanilla_sgd:
      context.sgd_weights = [
        w - context.learning_rate * grad
        for w, grad in zip(context.current_weights, context.gradients)
      ]
    return context

  @classmethod
  def _build_body(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question body and collect answers."""
    body = ca.Section()
    answers = []

    # Question description
    body.add_element(ca.Paragraph([
      "Consider the optimization problem of minimizing the function:"
    ]))

    body.add_element(ca.Equation(
      sp.latex(context.function),
      inline=False
    ))

    body.add_element(ca.Paragraph([
      "The gradient is:"
    ]))

    body.add_element(ca.Equation(
      f"\\nabla f = {sp.latex(context.gradient_function)}",
      inline=False
    ))

    # Current state
    body.add_element(ca.Paragraph([
      ca.Text("Current optimization state:", emphasis=True)
    ]))

    body.add_element(ca.Paragraph([
      "Current weights: ",
      ca.Equation(f"{format_vector(context.current_weights)}", inline=True)
    ]))

    body.add_element(ca.Paragraph([
      "Previous velocity: ",
      ca.Equation(f"{format_vector(context.prev_velocity)}", inline=True)
    ]))

    # Hyperparameters
    body.add_element(ca.Paragraph([
      ca.Text("Hyperparameters:", emphasis=True)
    ]))

    body.add_element(ca.Paragraph([
      "Learning rate: ",
      ca.Equation(f"\\alpha = {context.learning_rate}", inline=True)
    ]))

    body.add_element(ca.Paragraph([
      "Momentum coefficient: ",
      ca.Equation(f"\\beta = {context.momentum_beta}", inline=True)
    ]))

    # Questions
    body.add_element(ca.Paragraph([
      "Calculate the following updates:"
    ]))

    # Answer table
    table_data = []
    table_data.append(["Update Type", "Formula", "Result"])

    velocity_answer = ca.AnswerTypes.Vector(context.new_velocity, label="New velocity")
    weights_momentum_answer = ca.AnswerTypes.Vector(context.new_weights, label="Weights (momentum)")
    weights_sgd_answer = None

    answers.append(velocity_answer)
    table_data.append([
      "New velocity",
      ca.Equation(r"v' = \beta v + (1-\beta)\nabla f", inline=True),
      velocity_answer
    ])

    answers.append(weights_momentum_answer)
    table_data.append([
      "Weights (momentum)",
      ca.Equation(r"w' = w - \alpha v'", inline=True),
      weights_momentum_answer
    ])

    if context.show_vanilla_sgd:
      weights_sgd_answer = ca.AnswerTypes.Vector(context.sgd_weights, label="Weights (vanilla SGD)")
      answers.append(weights_sgd_answer)
      table_data.append([
        "Weights (vanilla SGD)",
        ca.Equation(r"w' = w - \alpha \nabla f", inline=True),
        weights_sgd_answer
      ])

    body.add_element(ca.Table(data=table_data))

    return body, answers

  @classmethod
  def _build_explanation(cls, context) -> tuple[ca.Section, list[ca.Answer]]:
    """Build question explanation."""
    explanation = ca.Section()

    explanation.add_element(ca.Paragraph([
      "Momentum helps gradient descent by accumulating a velocity vector in directions of "
      "consistent gradient, allowing faster convergence and reduced oscillation."
    ]))

    # Step 1: Calculate new velocity
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 1: Update velocity using momentum", emphasis=True)
    ]))

    explanation.add_element(ca.Paragraph([
      "The momentum update formula is:"
    ]))

    explanation.add_element(ca.Equation(
      f"v' = \\beta v + (1 - \\beta) \\nabla f",
      inline=False
    ))

    # Show calculation for each component
    digits = ca.Answer.DEFAULT_ROUNDING_DIGITS
    for i in range(context.num_variables):
      # Round all intermediate values to avoid floating point precision issues
      beta_times_v = round(context.momentum_beta * context.prev_velocity[i], digits)
      one_minus_beta = round(1 - context.momentum_beta, digits)
      one_minus_beta_times_grad = round((1 - context.momentum_beta) * context.gradients[i], digits)

      explanation.add_element(ca.Equation(
        f"v'[{i}] = {context.momentum_beta} \\times {context.prev_velocity[i]:.{digits}f} + "
        f"{one_minus_beta:.{digits}f} \\times {context.gradients[i]:.{digits}f} = "
        f"{beta_times_v:.{digits}f} + {one_minus_beta_times_grad:.{digits}f} = {context.new_velocity[i]:.{digits}f}",
        inline=False
      ))

    # Step 2: Update weights with momentum
    explanation.add_element(ca.Paragraph([
      ca.Text("Step 2: Update weights using new velocity", emphasis=True)
    ]))

    explanation.add_element(ca.Equation(
      f"w' = w - \\alpha v'",
      inline=False
    ))

    for i in range(context.num_variables):
      explanation.add_element(ca.Equation(
        f"w[{i}] = {context.current_weights[i]} - {context.learning_rate} \\times {context.new_velocity[i]:.4f} = {context.new_weights[i]:.4f}",
        inline=False
      ))

    # Comparison with vanilla SGD
    if context.show_vanilla_sgd:
      explanation.add_element(ca.Paragraph([
        ca.Text("Comparison with vanilla SGD:", emphasis=True)
      ]))

      explanation.add_element(ca.Paragraph([
        "Vanilla SGD (no momentum) would update directly using the gradient:"
      ]))

      explanation.add_element(ca.Equation(
        f"w' = w - \\alpha \\nabla f",
        inline=False
      ))

      for i in range(context.num_variables):
        explanation.add_element(ca.Equation(
          f"w[{i}] = {context.current_weights[i]} - {context.learning_rate} \\times {context.gradients[i]:.4f} = {context.sgd_weights[i]:.4f}",
          inline=False
        ))

      explanation.add_element(ca.Paragraph([
        "The momentum update differs because it incorporates the previous velocity, "
        "which can help accelerate learning and smooth out noisy gradients."
      ]))

    return explanation, []
