import math, re, torch
import torch.nn as nn
import torch.nn.functional as F

# =========================
# Model
# =========================
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.size()
        q = self.W_q(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
        mask = torch.tril(torch.ones(T, T, device=x.device))
        att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.dropout(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.W_o(y)

class GPTBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff=256, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, num_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.ReLU(),
            nn.Linear(d_ff, d_model), nn.Dropout(dropout)
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

class GPT(nn.Module):
    def __init__(self, vocab_size, d_model=128, num_heads=4, num_layers=2, max_len=128, dropout=0.1):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        self.blocks = nn.ModuleList([GPTBlock(d_model, num_heads, 256, dropout) for _ in range(num_layers)])
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
        x = self.token_emb(idx) + self.pos_emb(pos)
        for blk in self.blocks:
            x = blk(x)
        x = self.ln_f(x)
        return self.head(x)

# =========================
# Tokenization helpers
# =========================
def encode(s: str, stoi: dict, space_id: int):
    return [stoi.get(c, space_id) for c in s]

def decode(ids, itos: dict):
    return "".join(itos[i] for i in ids)

# =========================
# Greedy generation
# =========================
def greedy_until_eos(model, idx, itos, max_new_tokens=256, stop_text="<EOS>"):
    gen = idx
    for _ in range(max_new_tokens):
        idx_cond = gen[:, -model.pos_emb.num_embeddings:]
        logits = model(idx_cond)[:, -1, :]
        next_id = torch.argmax(logits, dim=-1, keepdim=True)
        gen = torch.cat([gen, next_id], dim=1)
        if stop_text in decode(gen[0].tolist(), itos):
            break
    return gen

# =========================
# Number parsing
# =========================
_UNITS = {
    "zero":0,"one":1,"two":2,"three":3,"four":4,"five":5,"six":6,"seven":7,"eight":8,"nine":9,
    "ten":10,"eleven":11,"twelve":12,"thirteen":13,"fourteen":14,"fifteen":15,"sixteen":16,
    "seventeen":17,"eighteen":18,"nineteen":19
}
_TENS = {"twenty":20,"thirty":30,"forty":40,"fifty":50,"sixty":60,"seventy":70,"eighty":80,"ninety":90}
_SCALES = {"hundred":100,"thousand":1000,"million":1_000_000}

def words_to_int(text: str):
    t = text.lower().replace("-", " ")
    toks = re.findall(r"[a-z]+|\d+", t)
    for tok in toks:
        if tok.isdigit():
            return int(tok)
    total, current, sign = 0, 0, 1
    for tok in toks:
        if tok in _UNITS: current += _UNITS[tok]
        elif tok in _TENS: current += _TENS[tok]
        elif tok == "hundred": current = (current or 1) * 100
        elif tok == "thousand": total += (current or 1) * 1000; current = 0
        elif tok == "million": total += (current or 1) * 1_000_000; current = 0
        elif tok == "minus" or tok == "negative": sign = -1
    return sign * (total + current)

def parse_number(piece: str):
    m = re.search(r"[-+]?\d[\d,]*\.?\d*", piece.replace(" ", ""))
    if m:
        s = m.group(0).replace(",", "")
        try:
            return int(s) if "." not in s else int(round(float(s)))
        except ValueError:
            return None
    return words_to_int(piece)

def extract_numbers(text: str):
    t = text.lower().strip()
    parts = re.split(r"[+]|(?:and)|,| ", t)
    nums = [parse_number(p) for p in parts if p.strip()]
    return [n for n in nums if n is not None]

# =========================
# Helpers for final answer
# =========================
def patch_final_answer(text: str, nums=None) -> str:
    try:
        s = sum(nums)
        text = re.sub(r"(=\s*)(-?\d+)", fr"\g<1>{s}", text, count=1)
        if "final answer:" in text:
            text = re.sub(r"(final answer:\s*)(-?\d+)", fr"\g<1>{s}", text)
        return text
    except Exception:
        return text

def render_add_steps(nums: list[int]) -> str:
    steps = []
    total = nums[0]
    steps.append(f"step 1: take {nums[0]}.")
    for i, n in enumerate(nums[1:], start=2):
        prev = total
        total += n
        steps.append(f"step {i}: add {n}. {prev} + {n} = {total}.")
    steps.append(f"final answer: {total}")
    return " ".join(steps)