#!/usr/bin/env python3
"""Checked standard-library walkthrough: 1-NN scaling and a Gini tree stump."""

from math import dist


TRAIN = [
    (850.0, 0.1, 0),
    (0.0, 0.8, 1),
    (1000.0, 0.7, 1),
]
QUERY = (900.0, 0.9)


def nearest(query, rows, scales=(1.0, 1.0)):
    query_scaled = tuple(value / scale for value, scale in zip(query, scales))
    ranked = []
    for amount, quote, label in rows:
        point_scaled = (amount / scales[0], quote / scales[1])
        ranked.append((dist(query_scaled, point_scaled), label))
    return min(ranked)


def gini(labels):
    if not labels:
        return 0.0
    proportions = [labels.count(label) / len(labels) for label in set(labels)]
    return 1.0 - sum(proportion**2 for proportion in proportions)


def best_quote_split(rows):
    values = sorted({row[1] for row in rows})
    thresholds = [(left + right) / 2 for left, right in zip(values, values[1:])]
    candidates = []
    for threshold in thresholds:
        left = [label for _, quote, label in rows if quote <= threshold]
        right = [label for _, quote, label in rows if quote > threshold]
        weighted = (len(left) * gini(left) + len(right) * gini(right)) / len(rows)
        candidates.append((weighted, threshold, left, right))
    return min(candidates)


def main():
    raw_distance, raw_label = nearest(QUERY, TRAIN)
    scaled_distance, scaled_label = nearest(QUERY, TRAIN, scales=(1000.0, 1.0))
    weighted_gini, threshold, left, right = best_quote_split(TRAIN)
    left_label = max(set(left), key=left.count)
    right_label = max(set(right), key=right.count)
    tree_label = left_label if QUERY[1] <= threshold else right_label

    print(f"KNN unscaled: label={raw_label}, distance={raw_distance:.3f}")
    print(f"KNN scaled: label={scaled_label}, distance={scaled_distance:.3f}")
    print(f"Best tree split: quote <= {threshold:.3f}, weighted_gini={weighted_gini:.3f}")
    print(f"Tree prediction: label={tree_label}")


if __name__ == "__main__":
    main()