pokemon-color-search/nearest.py

#!/usr/bin/env python3
import csv
import colorsys
import random
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument("color", nargs="?", default=None, help="Target color, randomized if not provided")
parser.add_argument("-n", "--number", default=1, type=int, help="Number of Pokemon to find")
parser.add_argument("-c", "--closeness", default=2, type=int, help="Closeness coefficient")
parser.add_argument("--database", default="database.csv", help="Database file")
parser.add_argument("-v", "--verbose", action="store_true", help="Print raw scores")
args = parser.parse_args()

if args.number <= 0:
    raise ValueError("Must request a number greater than 0")

if args.color is not None:
    if len(args.color) != 6:
        raise ValueError("Color must be a 6 digit hex")
    color = (int(args.color[0:2], base=16), int(args.color[2:4], base=16), int(args.color[4:6], base=16))
else:
    color = tuple(int(c * 255) for c in colorsys.hsv_to_rgb(random.random(), 0.9, 0.9))
    print(f"Generated random color: #{''.join(hex(c)[2:] for c in color)} / {color}")

results = []
with open(args.database) as infile:
    for name, x, *y in csv.reader(infile, delimiter=",", quotechar="'"):
        score = float(x) - args.closeness * sum(float(y_c) * c for y_c, c in zip(y, color))
        results.append((score, name))

if args.number > 1:
    for i, (score, name) in enumerate(sorted(results)[:args.number]):
        print(f"{i+1}:\t{name}")
        if args.verbose:
            print(f"\t\t\t{score=}")
else:
    best_score, best_name = min(results)
    print(f"Best match: {best_name}")
    if args.verbose:
        print(f"{score=}")