"""Train the laboratory's ten-digit linear reader, using arithmetic and NumPy.

From the repository root:
    python3 book/digit-reader-laboratory/building-and-running-our-reader/code/train_reader.py
Or download this file and optdigits-all.csv into the same folder and run:
    python3 train_reader.py --data optdigits-all.csv --output my-reader.json

No machine-learning library or automatic differentiation is used.
"""
import argparse
import csv
import json
from pathlib import Path
import time
import numpy as np


def read_data(path):
    groups = {"train": ([], []), "test": ([], [])}
    with open(path, newline="") as source:
        for row in csv.DictReader(source):
            images, labels = groups[row["split"]]
            images.append([1.0] + [int(row[f"p{j}"]) / 16 for j in range(64)])
            labels.append(int(row["label"]))
    return {name: (np.array(x, dtype=float), np.array(y, dtype=int))
            for name, (x, y) in groups.items()}


def scores_one(theta, image):
    """Ten dot products, with the constant bias coordinate already in image."""
    return [sum(weight * pixel for weight, pixel in zip(row, image))
            for row in theta]


def probabilities_one(scores):
    largest = max(scores)
    exponentials = [np.exp(score - largest) for score in scores]
    total = sum(exponentials)
    return [value / total for value in exponentials]


def one_image_gradient(theta, image, label):
    probabilities = probabilities_one(scores_one(theta, image))
    return [[(probabilities[k] - int(k == label)) * pixel for pixel in image]
            for k in range(10)]


def loss_and_gradient(theta, images, labels, lam=1.0, need_gradient=True):
    # Row i of scores contains the ten dot products for image i.
    scores = images @ theta.T
    largest = scores.max(axis=1, keepdims=True)
    exponentials = np.exp(scores - largest)
    totals = exponentials.sum(axis=1, keepdims=True)
    losses = largest[:, 0] + np.log(totals[:, 0]) - scores[np.arange(len(labels)), labels]
    loss = float(losses.sum() + lam * np.sum(theta * theta) / 2)
    if not need_gradient:
        return loss
    residuals = exponentials / totals
    residuals[np.arange(len(labels)), labels] -= 1
    # Entry (k,j) is sum_i residuals[i,k] * images[i,j].
    grad = residuals.T @ images + lam * theta
    return loss, grad


def check_gradient():
    """A small ten-class problem checks loops, bulk arithmetic and differences."""
    rng = np.random.default_rng(7)
    images = np.column_stack([np.ones(6), rng.random((6, 2))])
    labels = np.array([0, 3, 9, 3, 5, 1])
    theta = rng.normal(0, .2, (10, 3))
    _, grad = loss_and_gradient(theta, images, labels)
    explicit = theta.copy()  # lambda = 1, including the bias
    for image, label in zip(images, labels):
        explicit += np.array(one_image_gradient(theta, image, label))
    assert np.allclose(explicit, grad, atol=1e-12)
    worst = 0.0
    for index in np.ndindex(theta.shape):
        plus, minus = theta.copy(), theta.copy()
        plus[index] += 1e-5
        minus[index] -= 1e-5
        numeric = (loss_and_gradient(plus, images, labels, need_gradient=False)
                   - loss_and_gradient(minus, images, labels, need_gradient=False)) / 2e-5
        worst = max(worst, abs(numeric - grad[index]))
    assert worst < 1e-7, worst
    print(f"Derivative check: largest absolute error {worst:.3g}")
    return worst


def fit(images, labels, iterations=4000, lam=1.0):
    theta = np.zeros((10, images.shape[1]))
    history = []
    eta = .001
    for step in range(iterations + 1):
        loss, grad = loss_and_gradient(theta, images, labels, lam)
        norm = float(np.linalg.norm(grad))
        if step % 250 == 0 or step == iterations or norm < 1e-5:
            history.append({"iteration": step, "loss": loss, "gradientNorm": norm,
                            "theta": theta.ravel().tolist()})
            print(f"step {step:4d}: J={loss:.6f}, |gradient|={norm:.6g}", flush=True)
        if step == iterations or norm < 1e-5:
            return theta, history
        trial = min(2 * eta, .05)
        for _ in range(40):
            candidate = theta - trial * grad
            new_loss = loss_and_gradient(candidate, images, labels, lam, False)
            if new_loss <= loss - 1e-4 * trial * norm ** 2:
                break
            trial /= 2
        else:
            raise ArithmeticError("Step search stalled; inspect gradient and precision.")
        theta, eta = candidate, trial


def evaluate(theta, images, labels):
    scores = images @ theta.T
    predictions = scores.argmax(axis=1)
    confusion = np.zeros((10, 10), dtype=int)
    for actual, predicted in zip(labels, predictions):
        confusion[actual, predicted] += 1
    correct = int(np.sum(predictions == labels))
    return {"count": len(labels), "correct": correct, "accuracy": correct / len(labels),
            "confusion": confusion.tolist()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    local = Path(__file__).with_name("optdigits-all.csv")
    default = local if local.exists() else Path("demo-kit/shared/ml/data/optdigits-all.csv")
    parser.add_argument("--data", type=Path, default=default)
    parser.add_argument("--output", type=Path, default=Path("my-reader.json"))
    parser.add_argument("--iterations", type=int, default=4000)
    parser.add_argument("--check-only", action="store_true")
    args = parser.parse_args()
    error = check_gradient()
    if args.check_only:
        return
    if args.iterations < 0:
        parser.error("--iterations must be nonnegative")
    data = read_data(args.data)
    start = time.perf_counter()
    theta, history = fit(*data["train"], iterations=args.iterations)
    train = evaluate(theta, *data["train"])
    test = evaluate(theta, *data["test"])
    payload = {"schemaVersion": 1, "classes": list(range(10)), "pixelDivisor": 16,
               "biasFirst": True, "lambda": 1, "regularizesBias": True,
               "theta": theta.ravel().tolist(), "history": history,
               "train": train, "test": test, "gradientCheckError": error}
    args.output.write_text(json.dumps(payload, indent=2) + "\n")
    for name, result in [("train", train), ("test", test)]:
        print(f"{name}: {result['correct']}/{result['count']} = {result['accuracy']:.4%}")
    print(f"Saved {args.output} in {time.perf_counter()-start:.1f}s")


if __name__ == "__main__":
    main()