"""Colour classes of the user's picture (reference/user_picture.jpg, 827 x 803).

python tg_classes.py   -> reference/classes.png (flat colours per class), reference/classes.npy (class index per pixel)

The background is a smooth grey-blue gradient (darker at the top left); it is modelled by a cubic in x, y per channel, fitted
to the pixels that look like it, and every pixel near that model is background. The rest is sorted by hue, saturation and value:
"""
from pathlib import Path
import colorsys
import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
CLASSES = ["bg", "orange", "red", "bluegrey", "dark", "black", "pale", "navy"]
COLOURS = [(235, 235, 235), (240, 150, 40), (200, 30, 20), (110, 150, 175), (60, 60, 70), (0, 0, 0), (200, 240, 255),
           (30, 60, 160)]


def poly_terms(x, y):
    return np.stack([np.ones_like(x), x, y, x * x, x * y, y * y, x ** 3, x * x * y, x * y * y, y ** 3], -1)


def main():
    im = np.asarray(Image.open(ROOT / "reference" / "user_picture.jpg").convert("RGB")).astype(float) / 255.0
    H, W, _ = im.shape
    mx = im.max(2)
    mn = im.min(2)
    v = mx
    s = np.where(mx > 0, (mx - mn) / np.maximum(mx, 1e-6), 0)
    hsv = np.asarray(Image.fromarray((im * 255).astype(np.uint8)).convert("HSV")).astype(float)
    h = hsv[..., 0] / 255.0 * 360.0
    yy, xx = np.mgrid[0:H, 0:W].astype(float)
    xn, yn = xx / W, yy / H
    # background candidates: low saturation, light, blue-grey hue, outside the sub's rough box
    sub_box = (xx > 10) & (xx < 830) & (yy > 80) & (yy < 690)
    cand = (s < 0.26) & (v > 0.48) & (h > 165) & (h < 215) & ~sub_box
    T = poly_terms(xn[cand], yn[cand])
    coef = [np.linalg.lstsq(T, im[..., c][cand], rcond=None)[0] for c in range(3)]
    Tall = poly_terms(xn, yn)
    bg = np.stack([Tall @ coef[c] for c in range(3)], -1)
    d = np.sqrt(((im - bg) ** 2).sum(2))
    cls = np.zeros((H, W), np.uint8)
    is_bg = d < 0.075
    orange = (h >= 18) & (h <= 50) & (s > 0.38) & (v > 0.42)
    red = ((h < 18) | (h > 340)) & (s > 0.42) & (v > 0.14)
    black = v < 0.12
    pale = (s < 0.22) & (v > 0.70)
    navy = (h > 205) & (h < 240) & (s > 0.55) & (v > 0.18)
    bluegrey = (h > 180) & (h < 230) & (s > 0.12) & (v > 0.30)
    dark = ~is_bg
    cls[dark] = 4
    cls[bluegrey & ~is_bg] = 3
    cls[navy & ~is_bg] = 7
    cls[pale & ~is_bg] = 6
    cls[black] = 5
    cls[red & ~is_bg] = 2
    cls[orange & ~is_bg] = 1
    cls[is_bg] = 0
    np.save(ROOT / "reference" / "classes.npy", cls)
    np.save(ROOT / "reference" / "bg_model.npy", bg.astype(np.float32))
    out = np.array(COLOURS, np.uint8)[cls]
    Image.fromarray(out).save(ROOT / "reference" / "classes.png")
    for i, n in enumerate(CLASSES):
        print(n, int((cls == i).sum()))


if __name__ == "__main__":
    main()