"""
TAI 100, Week 5: Statistical learning. Generalisation and memorisation.

The task: decide whether a mango is ripe from two measurements, its colour
(0 = green, 1 = yellow) and its firmness (0 = soft, 1 = hard). The true rule
is a curved boundary, and 10% of the labels are wrong, because people
sorting mangoes make mistakes. So the best possible accuracy is about 0.90.

The learner: k nearest neighbours. To label a new mango, find the k most
similar mangoes in the training data and take a majority vote.

Part A  Training and test accuracy as k changes (memorising vs smoothing).
Part B  Shuffled labels: a learner can fit nonsense perfectly.
Part C  More data.
Part D  Choosing k fairly, with a validation set.

Run:  python3 week05_learning.py
Only the Python standard library is needed. Takes about a minute.
"""
import heapq
import random

LABEL_NOISE = 0.10


def truly_ripe(colour, firmness):
    return (1 - colour) ** 2 + firmness ** 2 < 0.5


def make_mangoes(n, rng):
    out = []
    for _ in range(n):
        c, f = rng.random(), rng.random()
        ripe = truly_ripe(c, f)
        if rng.random() < LABEL_NOISE:
            ripe = not ripe
        out.append((c, f, ripe))
    return out


def knn_predict(train, c, f, k):
    nearest = heapq.nsmallest(k, train, key=lambda m: (m[0] - c) ** 2 + (m[1] - f) ** 2)
    return 2 * sum(m[2] for m in nearest) > k


def accuracy(train, data, k):
    return sum(knn_predict(train, c, f, k) == ripe for c, f, ripe in data) / len(data)


def shuffled(data, rng):
    labels = [m[2] for m in data]
    rng.shuffle(labels)
    return [(c, f, r) for (c, f, _), r in zip(data, labels)]


if __name__ == "__main__":
    rng = random.Random(2026)
    test = make_mangoes(2_000, rng)
    trains = [make_mangoes(200, rng) for _ in range(10)]

    print("Part A: 200 training mangoes, averaged over 10 training sets")
    print("    k   training acc   test acc")
    for k in (1, 3, 7, 15, 31, 61, 121, 199):
        tr = sum(accuracy(t, t, k) for t in trains) / len(trains)
        te = sum(accuracy(t, test, k) for t in trains) / len(trains)
        print(f"  {k:>3}   {tr:>12.3f}   {te:>8.3f}")
    always = sum(not m[2] for m in test) / len(test)
    print(f"  always say 'not ripe': test acc {always:.3f}")
    perfect = sum(truly_ripe(c, f) == ripe for c, f, ripe in test) / len(test)
    print(f"  knowing the true rule:  test acc {perfect:.3f}  (the best possible on this test set)")

    print("\nPart B: the same learner, k = 1, on shuffled labels")
    t = trains[0]
    s = shuffled(t, rng)
    print(f"  true labels:     training acc {accuracy(t, t, 1):.3f}   test acc {accuracy(t, test, 1):.3f}")
    print(f"  shuffled labels: training acc {accuracy(s, s, 1):.3f}   test acc {accuracy(s, test, 1):.3f}")

    print("\nPart C: test accuracy against training set size (average of 5 sets)")
    print("      n    k = 1   k = 15")
    for n in (25, 50, 100, 200, 400, 800, 1600):
        sets = [make_mangoes(n, rng) for _ in range(5)]
        a1 = sum(accuracy(d, test, 1) for d in sets) / 5
        a15 = sum(accuracy(d, test, 15) for d in sets) / 5
        print(f"  {n:>5}   {a1:.3f}    {a15:.3f}")

    print("\nPart D: choosing k with a validation set (one training set of 200)")
    data = trains[0]
    fit, val = data[:150], data[150:]
    scores = {k: accuracy(fit, val, k) for k in (1, 3, 7, 15, 31, 61)}
    for k, a in scores.items():
        print(f"  k = {k:>2}: validation acc {a:.3f}")
    best = max(scores, key=scores.get)
    print(f"  chosen k = {best}; test acc using all 200 mangoes: {accuracy(data, test, best):.3f}")