"""Train the article's three-input, two-reader, one-output SMS classifier on CPU.

Requires PyTorch. Run: python train_sms.py
The first run downloads the UCI SMS Spam Collection into --data-dir.
Dataset: https://archive.ics.uci.edu/dataset/228/sms+spam+collection
Almeida & Hidalgo (2011), https://doi.org/10.24432/C5CC84, CC BY 4.0.
"""
import argparse
from collections import Counter
import hashlib
import io
import json
from pathlib import Path
import random
import re
import urllib.request
import zipfile

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

VOCABULARY = ("claim", "prize", "meeting")
DATA_URL = "https://archive.ics.uci.edu/static/public/228/sms%2Bspam%2Bcollection.zip"


class ReadWriteClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.W = nn.Linear(3, 2)  # Two rows: one incoming weight vector per hidden neuron.
        self.V = nn.Linear(2, 1)  # One row; one outgoing weight per hidden neuron.

    def forward(self, x):
        # x: [batch, 3]. Linear computes x @ W.weight.T + W.bias.
        h = F.relu(self.W(x))   # READ: [batch, 2] hidden activations.
        logits = self.V(h).squeeze(-1)  # WRITE + SUM: [batch] classification scores.
        return logits          # Positive favors spam; negative favors non-spam.


def encode(messages):
    rows = []
    for message in messages:
        counts = Counter(re.findall(r"[a-z]+", message.lower()))
        rows.append([counts[word] for word in VOCABULARY])
    return torch.tensor(rows, dtype=torch.float32)


def dataset(rows):
    return TensorDataset(
        encode([message for message, _ in rows]),
        torch.tensor([label for _, label in rows], dtype=torch.float32),
    )


def load_messages(data_dir):
    source = data_dir / "SMSSpamCollection"
    if not source.exists():
        data_dir.mkdir(parents=True, exist_ok=True)
        with urllib.request.urlopen(DATA_URL, timeout=60) as response:
            archive_bytes = response.read()
        with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive:
            # Read just the known file; do not extract arbitrary archive paths.
            source.write_bytes(archive.read("SMSSpamCollection"))
    raw = source.read_bytes()
    unique = {}
    total = 0
    for line in raw.decode("utf-8").splitlines():
        if not line.strip():
            continue
        raw_label, message = line.split("\t", 1)
        label = {"ham": 0, "spam": 1}[raw_label]
        # Keep duplicate messages out of both train and test, normalizing case
        # and whitespace. Different messages with identical counts stay distinct.
        identity = " ".join(message.lower().split())
        if identity in unique and unique[identity][1] != label:
            raise ValueError("A duplicate message has conflicting labels")
        unique.setdefault(identity, (message, label))
        total += 1
    return list(unique.values()), total, hashlib.sha256(raw).hexdigest()


def split_messages(rows, seed):
    # Stratify by class before constructing tensors; never fit on held-out rows.
    rng = random.Random(seed)
    train, test = [], []
    for label in (0, 1):
        group = [row for row in rows if row[1] == label]
        rng.shuffle(group)
        split = max(1, round(0.2 * len(group)))
        test.extend(group[:split])
        train.extend(group[split:])
    rng.shuffle(train)
    rng.shuffle(test)
    return train, test


@torch.no_grad()
def check_manual_example():
    model = ReadWriteClassifier()
    model.W.weight.copy_(torch.tensor([
        [1.5, 1.0, 0.0],
        [0.0, 0.0, 1.0],
    ]))
    model.V.weight.copy_(torch.tensor([[1.0, -1.0]]))
    model.W.bias.zero_()
    model.V.bias.zero_()
    x = encode([
        "Claim your prize. Claim it today.",
        "The project meeting is tomorrow.",
        "Prize!",
        "Claim your prize at the meeting.",
        "Prize meeting.",
        "Hello!",
    ])
    h = F.relu(model.W(x))
    torch.testing.assert_close(h, torch.tensor([
        [4.0, 0.0], [0.0, 1.0], [1.0, 0.0], [2.5, 1.0], [1.0, 1.0], [0.0, 0.0],
    ]))
    torch.testing.assert_close(model(x), torch.tensor([4.0, -1.0, 1.0, 1.5, 0.0, 0.0]))
    print("Manual arithmetic verified for six messages: 3 inputs → 2 readers → 1 score")


def train(model, train_loader, epochs, learning_rate):
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    loss_history = []
    for epoch in range(epochs):
        model.train()
        loss_sum = 0.0
        for word_counts, labels in train_loader:
            # 1. Reset: PyTorch accumulates gradients unless we clear them.
            optimizer.zero_grad()

            # 2. Forward: read with W, activate, write with V, add contributions.
            logits = model(word_counts)

            # 3. Loss: compare one raw score with each binary target (0 or 1).
            # This loss includes sigmoid internally; do not apply it first.
            loss = F.binary_cross_entropy_with_logits(logits, labels)

            # 4. Backward: calculate gradients for both matrices and biases.
            loss.backward()

            # 5. Update: Adam changes the parameters using those gradients.
            optimizer.step()
            loss_sum += loss.item() * len(labels)

        mean_loss = loss_sum / len(train_loader.dataset)
        loss_history.append(mean_loss)
        if epoch == 0 or (epoch + 1) % 10 == 0 or epoch + 1 == epochs:
            print(f"Epoch {epoch + 1:3d}: train cross-entropy = {mean_loss:.4f}")
    return loss_history


@torch.no_grad()
def evaluate(model, test_dataset):
    model.eval()
    x, labels = test_dataset.tensors
    logits = model(x)
    predicted = (logits > 0).long()  # Assign an exact zero to non-spam.
    # Rows are true classes; columns are predicted classes (non-spam, spam).
    confusion = torch.bincount(labels.long() * 2 + predicted, minlength=4).reshape(2, 2)
    tn, fp, fn, tp = confusion.flatten().tolist()
    return {
        "count": len(labels),
        "cross_entropy": F.binary_cross_entropy_with_logits(logits, labels).item(),
        "accuracy": (tn + tp) / len(labels),
        "always_nonspam_accuracy": (tn + fp) / len(labels),
        "spam_precision": tp / (tp + fp) if tp + fp else 0.0,
        "spam_recall": tp / (tp + fn) if tp + fn else 0.0,
        "confusion_true_rows_predicted_columns": confusion.tolist(),
        "zero_count_vectors": int((x.sum(dim=1) == 0).sum()),
    }


@torch.no_grad()
def explain(model, messages):
    model.eval()
    x = encode(messages)
    h = F.relu(model.W(x))
    # Each hidden activation scales its outgoing weight into the single score.
    contributions = h * model.V.weight[0]  # [batch, hidden neuron]
    logits = contributions.sum(dim=1) + model.V.bias[0]
    torch.testing.assert_close(logits, model(x))
    results = []
    for i, message in enumerate(messages):
        results.append({
            "message": message,
            "counts": x[i].tolist(),
            "hidden_activations": h[i].tolist(),
            "neuron_contributions": contributions[i].tolist(),
            "output_bias": model.V.bias[0].item(),
            "logit": logits[i].item(),
            "spam_probability": torch.sigmoid(logits[i]).item(),
            "prediction": "spam" if logits[i] > 0 else "non-spam",
        })
    return results


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=Path(".cache/read-write-sms"))
    parser.add_argument("--epochs", type=int, default=40)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--learning-rate", type=float, default=0.01)
    parser.add_argument("--output", type=Path, help="Optional JSON report path")
    args = parser.parse_args()
    if args.epochs < 1 or args.learning_rate <= 0:
        parser.error("epochs and learning-rate must be positive")
    torch.set_num_threads(1)
    check_manual_example()
    # Training starts from fresh random parameters, not the hand-chosen readers.
    torch.manual_seed(args.seed)
    rows, source_rows, digest = load_messages(args.data_dir)
    train_rows, test_rows = split_messages(rows, args.seed)
    train_dataset, test_dataset = dataset(train_rows), dataset(test_rows)
    train_loader = DataLoader(
        train_dataset, batch_size=64, shuffle=True,
        generator=torch.Generator().manual_seed(args.seed),
    )
    model = ReadWriteClassifier()
    print(f"Rows: {source_rows} original, {len(rows)} deduplicated; "
          f"{len(train_rows)} train, {len(test_rows)} held out")
    history = train(model, train_loader, args.epochs, args.learning_rate)
    metrics = evaluate(model, test_dataset)
    report = {
        "torch_version": torch.__version__, "seed": args.seed,
        "epochs": args.epochs, "learning_rate": args.learning_rate,
        "dataset_url": DATA_URL, "dataset_sha256": digest,
        "source_rows": source_rows, "unique_messages": len(rows),
        "train_count": len(train_rows), "vocabulary": VOCABULARY,
        "train_loss": history, "held_out": metrics,
        "parameters": {name: value.detach().tolist() for name, value in model.named_parameters()},
        "illustrative_messages": explain(model, [
            "Claim your prize. Claim it today.",
            "The project meeting is tomorrow.",
            "Hello!",
            "Prize!",
        ]),
    }
    print("Held-out metrics:", json.dumps(metrics, indent=2))
    print("One message, decomposed:", json.dumps(report["illustrative_messages"][0], indent=2))
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()