from __future__ import annotations
import numpy as np
from PIL import Image
from .common import to_rgb


Q_BITS = 5
Q_LEVELS = 1 << Q_BITS          # 32
LOG2_B = 3 * Q_BITS             # 15


def compute_colorfulness_hs_from_pil(pil: Image.Image) -> float:
    rgb = to_rgb(pil)
    arr = np.asarray(rgb, dtype=np.uint8)
    R = arr[..., 0].ravel()
    G = arr[..., 1].ravel()
    B = arr[..., 2].ravel()

    Rf = R.astype(np.float32); Gf = G.astype(np.float32); Bf = B.astype(np.float32)
    rg = Rf - Gf
    yb = (Rf + Gf) * 0.5 - Bf
    std_rg = float(np.std(rg, ddof=0))
    std_yb = float(np.std(yb, ddof=0))
    mean_rg = float(np.mean(rg))
    mean_yb = float(np.mean(yb))
    return round((std_rg**2 + std_yb**2) ** 0.5 + 0.3 * (mean_rg**2 + mean_yb**2) ** 0.5, 2)


def compute_color_richness_pct_from_pil(pil: Image.Image) -> float:
    rgb = to_rgb(pil)

    arr = np.asarray(rgb, dtype=np.uint8)
    rq = (arr[..., 0] >> (8 - Q_BITS)).astype(np.int32)
    gq = (arr[..., 1] >> (8 - Q_BITS)).astype(np.int32)
    bq = (arr[..., 2] >> (8 - Q_BITS)).astype(np.int32)
    key = (rq << (2 * Q_BITS)) | (gq << Q_BITS) | bq
    k = key.ravel()

    total = int(k.size) or 1
    hist = np.bincount(k, minlength=Q_LEVELS ** 3)
    probs = hist[hist > 0].astype(np.float64) / float(total)
    H = float(-(probs * np.log2(probs)).sum()) if probs.size else 0.0
    return float(round(100.0 * (H / LOG2_B), 1))


def compute_color_metrics_from_pil(pil: Image.Image) -> tuple[float, float]:
    cf = compute_colorfulness_hs_from_pil(pil)
    cr = compute_color_richness_pct_from_pil(pil)
    return cf, cr
