# Student copy for HW7, generated by build_student_modules.py from the course's instructor copy.
# Do not edit; a new copy may be posted if the course files change.
"""Shared Spanish -> English translation pipeline (CS6140 HW6 Problem 4; HW7 Problems 3-4).

HW6 and HW7 import this file unchanged, so every translation model in the course is trained on the same
pairs, uses the same vocabularies and is scored on the same test sentences with the same BLEU.

Data. Tatoeba sentence pairs (https://tatoeba.org, CC BY 2.0 FR), the course file
`data/sentence_pairs_large.tsv` (columns: id, Spanish, id, English). Text is lower-cased and
punctuation (including the Spanish marks the d2l book's French-English code does not know) becomes
separate tokens. Only pairs whose two sides both have at most `max_len` tokens are kept.

Splits are made by *Spanish sentence*, not by pair: a Spanish sentence with several English
translations goes entirely into one split. So a test sentence was never seen in training, and all its
translations serve as references for BLEU. The final-test split is never scored in HW6 or HW7; it is kept
untouched for a last comparison.

The data module subclasses the book's `d2l.MTFraEng` (d2l §10.5) and returns its batch format:
(src, tgt_input, src_valid_len, label), with src = tokens + <eos> + padding, tgt_input = <bos> + tokens,
label = tokens + <eos> (both padded), all of length num_steps = max_len + 1.

Models follow the book's encoder-decoder interface (d2l §10.6): `model(src, tgt_input, src_valid_len)`
returns logits (batch, steps, vocab); `model.encoder(src, src_valid_len)` and
`model.decoder.init_state(enc_outputs, src_valid_len)` / `model.decoder(Y, state)` decode step by step
(this is what `d2l.EncoderDecoder.predict_step` and `translate_greedy` below call).

Answer-only code. Blocks between `# --- SOLUTION-ONLY (begin)` and `# --- SOLUTION-ONLY (end)` contain the
answers to HW6 TODO blanks (masked loss, greedy decoding, the clipped n-gram count, BLEU by length). They are
stripped from the copy shipped with the HW6 starter; there the notebook passes its own functions in
(`loss_fn=`, `clip_fn=`, `bleu_by_len=`). HW7 ships the full module.

BLEU is the standard corpus BLEU-4 (Papineni et al., 2002), as reported by sacrebleu: clipped n-gram counts
over the whole test set, all references per sentence, brevity penalty against the closest reference
length, x100. (The book's `d2l.bleu` is a different, sentence-level variant; its numbers are not
comparable.)
"""
import collections
import json
import math
import os
import random
import re
import time

import torch
from torch import nn
import torch.nn.functional as F
from d2l import torch as d2l

BUCKETS = ((1, 4), (5, 6), (7, 8), (9, 12))     # Spanish sentence length (tokens) for BLEU by length


def tokenize(sentence):
    """Lower-case words (with accents and inner apostrophes) and punctuation, each a token."""
    sentence = sentence.lower().replace(" ", " ").replace("\xa0", " ").replace("’", "'")
    return re.findall(r"[\w']+|[.,!?¿¡;:\"]", sentence)


class TatoebaEsEn(d2l.MTFraEng):
    """Spanish -> English Tatoeba pairs as a d2l data module (see the module docstring)."""

    def __init__(self, root, batch_size=128, max_len=12, num_train=100_000, num_val=1_000, num_test=2_000,
                 num_final_test=2_000, min_freq=2, seed=0):
        d2l.DataModule.__init__(self)
        self.save_hyperparameters()
        self.num_steps = max_len + 1
        src, tgt = self._tokenize(self._preprocess(self._download()))
        by_src = collections.defaultdict(list)
        for es, en in zip(src, tgt):
            by_src[tuple(es)].append(en)
        sources = list(by_src)
        rng = random.Random(seed)
        rng.shuffle(sources)
        cut = [num_test, num_test + num_final_test, num_test + num_final_test + num_val]
        test, final, val, train = sources[:cut[0]], sources[cut[0]:cut[1]], sources[cut[1]:cut[2]], sources[cut[2]:]
        train_pairs = [(list(es), en) for es in train for en in by_src[es]]
        rng.shuffle(train_pairs)
        self.num_available = len(train_pairs)
        train_pairs = train_pairs[:num_train]
        self.arrays, self.src_vocab, self.tgt_vocab = self._build_arrays(*zip(*train_pairs))
        val_pairs = [(list(es), en) for es in val for en in by_src[es]]
        self.val_arrays = self._build_arrays(*zip(*val_pairs), self.src_vocab, self.tgt_vocab)[0]
        strip = lambda s: [w for w in s if w != "<eos>"]
        self.test_src = [strip(es) for es in test]
        self.test_refs = [[strip(en) for en in by_src[es]] for es in test]
        self.final_test_src = [strip(es) for es in final]
        self.final_test_refs = [[strip(en) for en in by_src[es]] for es in final]
        self.num_pairs, self.num_sources = len(src), len(sources)

    def _download(self):
        path = os.path.join(self.root, "sentence_pairs_large.tsv")
        lines = []
        with open(path, encoding="utf-8-sig") as f:
            for line in f:
                parts = line.rstrip("\n").split("\t")
                if len(parts) >= 4:
                    lines.append(parts[1].replace("\t", " ") + "\t" + parts[3].replace("\t", " "))
        return "\n".join(lines)

    def _preprocess(self, text):
        return text                                 # tokenize() does the lower-casing and splitting

    def _tokenize(self, text, max_examples=None):
        """(src, tgt) token lists, each ending in '<eos>'; pairs longer than max_len on either side dropped."""
        src, tgt = [], []
        for line in text.split("\n"):
            es, en = (tokenize(part) for part in line.split("\t"))
            if 0 < len(es) <= self.max_len and 0 < len(en) <= self.max_len:
                src.append(es + ["<eos>"])
                tgt.append(en + ["<eos>"])
        return src, tgt

    def _build_arrays(self, src, tgt, src_vocab=None, tgt_vocab=None):
        """The book's _build_arrays, but on token lists, with the vocabularies built from training data only."""
        def _build_array(sentences, vocab, is_tgt=False):
            pad_or_trim = lambda seq, t: seq[:t] if len(seq) > t else seq + ["<pad>"] * (t - len(seq))
            sentences = [pad_or_trim(list(s), self.num_steps) for s in sentences]
            if is_tgt:
                sentences = [["<bos>"] + s for s in sentences]
            if vocab is None:
                vocab = d2l.Vocab(sentences, min_freq=self.min_freq, reserved_tokens=["<pad>", "<bos>", "<eos>"])
            array = torch.tensor([vocab[s] for s in sentences])
            valid_len = (array != vocab["<pad>"]).type(torch.int32).sum(1)
            return array, vocab, valid_len
        src_array, src_vocab, src_valid_len = _build_array(src, src_vocab)
        tgt_array, tgt_vocab, _ = _build_array(tgt, tgt_vocab, True)
        return (src_array, tgt_array[:, :-1], src_valid_len, tgt_array[:, 1:]), src_vocab, tgt_vocab

    def get_dataloader(self, train):
        return self.get_tensorloader(self.arrays if train else self.val_arrays, train)

    def src_batch(self, sentences):
        """Token lists (no <eos>) -> (src, src_valid_len) tensors in the training format."""
        rows = [list(s)[:self.max_len] + ["<eos>"] for s in sentences]
        rows = [r + ["<pad>"] * (self.num_steps - len(r)) for r in rows]
        arr = torch.tensor([self.src_vocab[r] for r in rows])
        valid_len = (arr != self.src_vocab["<pad>"]).type(torch.int32).sum(1)
        return arr, valid_len

    def describe(self):
        n = len(self.arrays[0])
        return (f"Tatoeba es->en, both sides <= {self.max_len} tokens: {self.num_pairs:,} pairs, "
                f"{self.num_sources:,} distinct Spanish sentences.\n"
                f"train {n:,} pairs (of {self.num_available:,} available) | val {len(self.val_arrays[0]):,} pairs | "
                f"test {len(self.test_src):,} Spanish sentences "
                f"({sum(map(len, self.test_refs)) / len(self.test_refs):.2f} references each) | "
                f"final test {len(self.final_test_src):,} (kept untouched)\n"
                f"vocabularies: Spanish {len(self.src_vocab):,}, English {len(self.tgt_vocab):,} tokens "
                f"(min frequency {self.min_freq}); num_steps = {self.num_steps}")


# ------------------------------------------------------------------------------------------------
# The book's RNN decoder (d2l §10.7) -- the d2l *package* ships the encoder and d2l.Seq2Seq but not this
# class, nor the masked loss. Copied from the book (CC BY-SA 4.0), with nn.LazyLinear -> nn.Linear.
# ------------------------------------------------------------------------------------------------
class Seq2SeqDecoder(d2l.Decoder):
    """The book's RNN decoder: the encoder's last-step output is the context at every step."""

    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout=0):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = d2l.GRU(embed_size + num_hiddens, num_hiddens, num_layers, dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)
        self.apply(d2l.init_seq2seq)

    def init_state(self, enc_all_outputs, *args):
        return enc_all_outputs

    def forward(self, X, state):
        embs = self.embedding(X.t().type(torch.int32))         # (num_steps, batch, embed)
        enc_output, hidden_state = state
        context = enc_output[-1]                                # the encoder's output at its LAST position
        context = context.repeat(embs.shape[0], 1, 1)
        embs_and_context = torch.cat((embs, context), -1)
        outputs, hidden_state = self.rnn(embs_and_context, hidden_state)
        outputs = self.dense(outputs).swapaxes(0, 1)            # (batch, num_steps, vocab)
        return outputs, [enc_output, hidden_state]


class BookSeq2Seq(d2l.Seq2Seq):
    """d2l.Seq2Seq plus the book's masked loss (padding positions do not count)."""

    def loss(self, Y_hat, Y):
        l = d2l.Classifier.loss(self, Y_hat, Y, averaged=False)
        mask = (Y.reshape(-1) != self.tgt_pad).type(torch.float32)
        return (l * mask).sum() / mask.sum()


def book_seq2seq(data, embed_size=256, num_hiddens=512, num_layers=2, dropout=0.2, lr=1e-3):
    """The book's package baseline: d2l.Seq2SeqEncoder + the book's decoder, wrapped in d2l.Seq2Seq."""
    encoder = d2l.Seq2SeqEncoder(len(data.src_vocab), embed_size, num_hiddens, num_layers, dropout)
    decoder = Seq2SeqDecoder(len(data.tgt_vocab), embed_size, num_hiddens, num_layers, dropout)
    return BookSeq2Seq(encoder, decoder, tgt_pad=data.tgt_vocab["<pad>"], lr=lr)


def packed_rnn(rnn, embs, valid_len):
    """Run a time-major nn.GRU / nn.LSTM over padded input, *skipping the padding*.

    embs: (num_steps, batch, embed); valid_len: (batch,). Returns outputs (num_steps, batch, hidden), zero
    at padded positions, and the final state taken at each sentence's own last token (not after the padding)."""
    packed = nn.utils.rnn.pack_padded_sequence(embs, valid_len.cpu(), enforce_sorted=False)
    outputs, state = rnn(packed)
    outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs, total_length=embs.shape[0])
    return outputs, state


# ------------------------------------------------------------------------------------------------
# Training (provided plumbing)
# ------------------------------------------------------------------------------------------------
# --- SOLUTION-ONLY (begin): HW6 P4(B) masked_ce; without it, pass loss_fn to train_seq2seq ---
def masked_loss(logits, labels, pad):
    """Cross-entropy averaged over the non-padding target positions."""
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1), ignore_index=pad)
# --- SOLUTION-ONLY (end) ---


@torch.no_grad()
def eval_loss(model, data, device, loss_fn, batch_size=512):
    model.eval()
    tot, n = 0.0, 0
    pad = data.tgt_vocab["<pad>"]
    arrays = data.val_arrays
    for i in range(0, len(arrays[0]), batch_size):
        src, tgt_in, src_len, label = (a[i:i + batch_size].to(device) for a in arrays)
        k = (label != pad).sum().item()               # loss_fn averages over real tokens: weight by their count
        tot += loss_fn(model(src, tgt_in, src_len), label).item() * k
        n += k
    model.train()
    return tot / n


def train_seq2seq(model, data, device, steps, loss_fn=None, lr=1e-3, clip=1.0, progress=None, desc="seq2seq",
                  seed=0):
    """AdamW on random training batches + gradient clipping; returns a history dict.
    Every run with the same seed sees the same sequence of batches."""
    pad = data.tgt_vocab["<pad>"]
    if loss_fn is None:                    # (masked_loss is not in the HW6 student copy: pass loss_fn there)
        loss_fn = lambda logits, labels: masked_loss(logits, labels, pad)
    model.to(device).train()
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    g = torch.Generator().manual_seed(seed)
    n = len(data.arrays[0])
    eval_every = max(steps // 10, 1)
    hist = {"step": [], "train_loss": [], "val_loss": []}
    bar = progress(steps, desc) if progress else None
    running = None
    t0 = time.time()
    for step in range(1, steps + 1):
        idx = torch.randint(n, (data.batch_size,), generator=g)
        src, tgt_in, src_len, label = (a[idx].to(device) for a in data.arrays)
        loss = loss_fn(model(src, tgt_in, src_len), label)
        opt.zero_grad()
        loss.backward()
        if clip:
            nn.utils.clip_grad_norm_(model.parameters(), clip)
        opt.step()
        running = loss.item() if running is None else 0.98 * running + 0.02 * loss.item()
        if bar is not None:
            bar.update(1)
        if step % eval_every == 0 or step == steps:
            hist["step"].append(step)
            hist["train_loss"].append(running)
            hist["val_loss"].append(eval_loss(model, data, device, loss_fn))
            if bar is not None:
                bar.set_postfix(val_loss=f"{hist['val_loss'][-1]:.3f}")
    if bar is not None:
        bar.close()
    hist["seconds"] = time.time() - t0
    hist["steps"] = steps
    return hist


# ------------------------------------------------------------------------------------------------
# Decoding
# ------------------------------------------------------------------------------------------------
def ids_to_tokens(ids, vocab):
    """One output row of token IDs -> tokens, cut at the first <eos> (and never containing <pad>)."""
    eos, pad = vocab["<eos>"], vocab["<pad>"]
    out = []
    for i in ids:
        if i in (eos, pad):
            break
        out.append(vocab.idx_to_token[i])
    return out


# --- SOLUTION-ONLY (begin): HW6 P4(C) greedy_translate ---
@torch.no_grad()
def translate_greedy(model, data, sentences, device, max_steps=None, batch_size=500):
    """Greedy decoding through the book's encoder-decoder interface (like d2l's predict_step, batched).
    sentences: token lists. Returns token lists without <eos>."""
    model.eval()
    max_steps = max_steps or data.num_steps + 2
    bos = data.tgt_vocab["<bos>"]
    hyps = []
    for i in range(0, len(sentences), batch_size):
        src, src_len = (a.to(device) for a in data.src_batch(sentences[i:i + batch_size]))
        state = model.decoder.init_state(model.encoder(src, src_len), src_len)
        Y = torch.full((len(src), 1), bos, device=device)
        out = []
        for _ in range(max_steps):
            logits, state = model.decoder(Y, state)
            Y = logits[:, -1:].argmax(-1)
            out.append(Y)
        hyps += [ids_to_tokens(row, data.tgt_vocab) for row in torch.cat(out, 1).tolist()]
    model.train()
    return hyps
# --- SOLUTION-ONLY (end) ---


# ------------------------------------------------------------------------------------------------
# BLEU
# ------------------------------------------------------------------------------------------------
def ngrams(tokens, n):
    return collections.Counter(tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1))


# --- SOLUTION-ONLY (begin): HW6 P4(D) clipped_matches ---
def clipped_matches(hyp, refs, n):
    """n-grams of hyp that occur in a reference, each counted at most as often as in any ONE reference."""
    most = collections.Counter()
    for r in refs:
        for g, c in ngrams(r, n).items():
            most[g] = max(most[g], c)
    return sum(min(c, most[g]) for g, c in ngrams(hyp, n).items())
# --- SOLUTION-ONLY (end) ---


def _clip_fn(clip_fn):
    if clip_fn is not None:
        return clip_fn
    if "clipped_matches" not in globals():
        raise RuntimeError("this copy of seq2seq_common has no clipped_matches: pass clip_fn=your clipped_matches")
    return globals()["clipped_matches"]


def corpus_bleu(hyps, refs_list, max_n=4, clip_fn=None):
    """Standard corpus BLEU-4 (x100). hyps: token lists; refs_list: for each hyp, a list of reference token lists.
    clip_fn(hyp, refs, n): the clipped n-gram count (default: this module's clipped_matches)."""
    clip = _clip_fn(clip_fn)
    matches, totals = [0] * max_n, [0] * max_n
    hyp_len = ref_len = 0
    for hyp, refs in zip(hyps, refs_list):
        hyp_len += len(hyp)
        ref_len += min((abs(len(r) - len(hyp)), len(r)) for r in refs)[1]       # closest reference length
        for n in range(1, max_n + 1):
            matches[n - 1] += clip(hyp, refs, n)                                # clipped counts
            totals[n - 1] += max(len(hyp) - n + 1, 0)
    if min(matches) == 0:
        return 0.0
    log_p = sum(math.log(m / t) for m, t in zip(matches, totals)) / max_n
    bp = 1.0 if hyp_len > ref_len else math.exp(1 - ref_len / max(hyp_len, 1))
    return 100 * bp * math.exp(log_p)


# --- SOLUTION-ONLY (begin): HW6 P4(D) bleu_by_length; without it, pass bleu_by_len to save_results ---
def bleu_by_length(hyps, refs_list, srcs, buckets=BUCKETS, clip_fn=None):
    """{'1-4': BLEU, ...} over the test sentences whose SOURCE length (in tokens) falls in each bucket."""
    out = {}
    for lo, hi in buckets:
        keep = [i for i, s in enumerate(srcs) if lo <= len(s) <= hi]
        out[f"{lo}-{hi}"] = corpus_bleu([hyps[i] for i in keep], [refs_list[i] for i in keep], clip_fn=clip_fn)
    return out
# --- SOLUTION-ONLY (end) ---


# ------------------------------------------------------------------------------------------------
# Save / load (HW6 writes these; HW7 reads them)
# ------------------------------------------------------------------------------------------------
def save_results(path, name, data, hyps, hist=None, extra=None, clip_fn=None, bleu_by_len=None):
    """Write results_hw6.json-style results: BLEU overall and by length, plus every test translation.
    clip_fn: passed to corpus_bleu; bleu_by_len: precomputed {'1-4': BLEU, ...} (else computed here)."""
    if bleu_by_len is None:
        bleu_by_len = bleu_by_length(hyps, data.test_refs, data.test_src, clip_fn=clip_fn)
    res = {
        "model": name,
        "bleu": round(corpus_bleu(hyps, data.test_refs, clip_fn=clip_fn), 2),
        "bleu_by_src_len": {k: round(v, 2) for k, v in bleu_by_len.items()},
        "data": {k: getattr(data, k) for k in ("max_len", "num_train", "num_val", "num_test", "min_freq", "seed")},
        "test_src": [" ".join(s) for s in data.test_src],
        "test_hyp": [" ".join(h) for h in hyps],
    }
    if hist is not None:
        res["train"] = {"steps": hist["steps"], "seconds": round(hist["seconds"], 1),
                        "final_val_loss": round(hist["val_loss"][-1], 4)}
    res.update(extra or {})
    with open(path, "w", encoding="utf-8") as f:
        json.dump(res, f, indent=1, ensure_ascii=False)
    return res


def load_results(path):
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def save_checkpoint(path, model, config, data):
    """Weights (on CPU) + the constructor arguments + both vocabularies, so the model can be rebuilt."""
    torch.save({"state_dict": {k: v.cpu() for k, v in model.state_dict().items()}, "config": config,
                "src_vocab": data.src_vocab.idx_to_token, "tgt_vocab": data.tgt_vocab.idx_to_token}, path)


def load_checkpoint(path, model=None, data=None):
    """Returns the checkpoint dict; if `model` is given, loads the weights into it. If `data` is given,
    checks that its vocabularies are the ones the checkpoint was trained with."""
    ckpt = torch.load(path, map_location="cpu")
    if data is not None:
        assert ckpt["src_vocab"] == data.src_vocab.idx_to_token and ckpt["tgt_vocab"] == data.tgt_vocab.idx_to_token, \
            "the checkpoint was trained with different vocabularies (different data settings?)"
    if model is not None:
        model.load_state_dict(ckpt["state_dict"])
    return ckpt