import logging
from typing import List, Tuple

import numpy as np

import QuizGenerator.contentast as ca
from QuizGenerator.mixins import TableQuestionMixin
from QuizGenerator.question import QuestionRegistry

from .matrices import MatrixQuestion

log = logging.getLogger(__name__)


@QuestionRegistry.register("cst463.attention.forward-pass")
class AttentionForwardPass(MatrixQuestion, TableQuestionMixin):
  
  @staticmethod
  def simple_attention(Q, K, V):
    """
    Q: (seq_len, d_k) - queries
    K: (seq_len, d_k) - keys
    V: (seq_len, d_v) - values

    Returns: (seq_len, d_v) - attended output
    """
    d_k = Q.shape[1]
    
    # Compute attention scores
    scores = Q @ K.T / np.sqrt(d_k)
    
    # Softmax to get weights
    attention_weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
    
    # Weighted sum of values
    output = attention_weights @ V
    
    return output, attention_weights
  
  @classmethod
  def _build_context(cls, *, rng_seed=None, **kwargs):
    rng = cls.get_rng(rng_seed)
    digits = cls.get_digits_to_round(digits_to_round=kwargs.get("digits_to_round"))

    seq_len = kwargs.get("seq_len", 3)
    d_k = kwargs.get("key_dimension", 1)
    d_v = kwargs.get("value_dimension", 1)

    Q = cls.get_rounded_matrix(rng, (seq_len, d_k), 0, 3, digits)
    K = cls.get_rounded_matrix(rng, (seq_len, d_k), 0, 3, digits)
    V = cls.get_rounded_matrix(rng, (seq_len, d_v), 0, 3, digits)

    output, weights = cls.simple_attention(Q, K, V)

    return {
      "digits": digits,
      "seq_len": seq_len,
      "d_k": d_k,
      "d_v": d_v,
      "Q": Q,
      "K": K,
      "V": V,
      "output": output,
      "weights": weights,
    }

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

    body.add_element(
      ca.Text("Given the below information about a self attention layer, please calculate the output sequence.")
    )
    body.add_element(
      cls.create_info_table(
        {
          "Q": ca.Matrix(context["Q"]),
          "K": ca.Matrix(context["K"]),
          "V": ca.Matrix(context["V"]),
        }
      )
    )

    weights_answer = ca.AnswerTypes.Matrix(context["weights"], label="Weights")
    output_answer = ca.AnswerTypes.Matrix(context["output"], label="Output")
    answers.append(weights_answer)
    answers.append(output_answer)
    body.add_elements([
      ca.LineBreak(),
      weights_answer,
      ca.LineBreak(),
      output_answer,
    ])

    return body, answers

  @classmethod
  def _build_explanation(cls, context) -> Tuple[ca.Section, List[ca.Answer]]:
    """Build question explanation."""
    explanation = ca.Section()
    digits = ca.Answer.DEFAULT_ROUNDING_DIGITS

    explanation.add_element(
      ca.Paragraph([
        "Self-attention uses scaled dot-product attention to compute a weighted combination of values based on query-key similarity."
      ])
    )

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

    d_k = context["Q"].shape[1]
    explanation.add_element(
      ca.Equation(f"\\text{{scores}} = \\frac{{Q K^T}}{{\\sqrt{{d_k}}}} = \\frac{{Q K^T}}{{\\sqrt{{{d_k}}}}}")
    )

    scores = context["Q"] @ context["K"].T / np.sqrt(d_k)

    explanation.add_element(
      ca.Paragraph([
        "Raw scores (scaling by ",
        ca.Equation(f'\\sqrt{{{d_k}}}', inline=True),
        " prevents extremely large values):"
      ])
    )
    explanation.add_element(ca.Matrix(np.round(scores, digits)))

    # Step 2: Apply softmax
    explanation.add_element(
      ca.Paragraph([
        ca.Text("Step 2: Apply softmax to get attention weights", emphasis=True)
      ])
    )

    explanation.add_element(
      ca.Equation(r"\alpha_{ij} = \frac{\exp(\text{score}_{ij})}{\sum_k \exp(\text{score}_{ik})}")
    )

    # Show ONE example row
    explanation.add_element(
      ca.Paragraph([
        "Example: Row 0 softmax computation"
      ])
    )

    row_scores = scores[0]
    exp_scores = np.exp(row_scores)
    sum_exp = exp_scores.sum()
    weights_row = exp_scores / sum_exp

    exp_terms = " + ".join([f"e^{{{s:.{digits}f}}}" for s in row_scores])

    explanation.add_element(
      ca.Paragraph([
        f"Denominator = {exp_terms} = {sum_exp:.{digits}f}"
      ])
    )

    # Format array with proper rounding
    weights_str = "[" + ", ".join([f"{w:.{digits}f}" for w in weights_row]) + "]"
    explanation.add_element(
      ca.Paragraph([
        f"Resulting weights: {weights_str}"
      ])
    )

    explanation.add_element(
      ca.Paragraph([
        "Complete attention weight matrix:"
      ])
    )
    explanation.add_element(ca.Matrix(np.round(context["weights"], digits)))

    # Step 3: Weighted sum of values
    explanation.add_element(
      ca.Paragraph([
        ca.Text("Step 3: Compute weighted sum of values", emphasis=True)
      ])
    )

    explanation.add_element(
      ca.Equation(r"\text{output} = \text{weights} \times V")
    )

    explanation.add_element(
      ca.Paragraph([
        "Final output:"
      ])
    )
    explanation.add_element(ca.Matrix(np.round(context["output"], digits)))

    return explanation, []
