"""
Independent color-space comparison: how much of the uploaded photo's color
distribution is also present in each template, plus which specific colors
matched. Entirely separate from the SIFT/ORB/SuperGlue/LoFTR scores and the
weighted final match -- nothing here feeds into those, it's its own section.

Template-side color data (histogram + dominant colors) is precomputed once
at startup from the same masked bgr/mask already produced during bootstrap,
same principle as the other methods: a request only ever analyzes the one
uploaded image.
"""

import cv2
import numpy as np

N_DOMINANT_COLORS = 5
HIST_HUE_BINS = 30
HIST_SAT_BINS = 32

# k-means on every foreground pixel is wasteful and, given this app's
# history of memory blowups, worth bounding explicitly -- a random sample
# is statistically equivalent for "what are the dominant colors" purposes.
MAX_KMEANS_SAMPLE = 20000

# Empirical ceiling for CIE76 Lab distance beyond which two colors are
# considered maximally dissimilar (similarity floors at 0%). ~100 comfortably
# covers the largest perceptual differences (e.g. black vs. white is ~100).
MAX_LAB_DISTANCE = 100.0

_template_color_data = {}  # name -> {"hist": ndarray, "colors": [...]}


def _masked_pixels_bgr(bgr, mask):
    ys, xs = np.where(mask > 0)
    if len(ys) == 0:
        return np.zeros((0, 3), dtype=np.uint8)
    return bgr[ys, xs]


def compute_hs_histogram(bgr, mask):
    """Hue+Saturation 2D histogram (Value/brightness deliberately excluded
    so lighting differences between photos don't masquerade as color
    differences), normalized so it sums to 1."""
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    hist = cv2.calcHist([hsv], [0, 1], mask, [HIST_HUE_BINS, HIST_SAT_BINS],
                        [0, 180, 0, 256])
    total = hist.sum()
    if total > 0:
        hist = hist / total
    return hist


def histogram_match_pct(hist_a, hist_b):
    """% of one distribution's mass that the other also covers -- both
    histograms are pre-normalized, so HISTCMP_INTERSECT directly sums
    min(a[i], b[i]) into a 0-1 overlap fraction."""
    intersection = cv2.compareHist(hist_a.astype(np.float32), hist_b.astype(np.float32),
                                    cv2.HISTCMP_INTERSECT)
    return round(float(intersection) * 100, 1)


def dominant_colors(bgr, mask, k=N_DOMINANT_COLORS):
    pixels = _masked_pixels_bgr(bgr, mask)
    if len(pixels) == 0:
        return []

    if len(pixels) > MAX_KMEANS_SAMPLE:
        idx = np.random.choice(len(pixels), MAX_KMEANS_SAMPLE, replace=False)
        pixels = pixels[idx]

    unique_count = len(np.unique(pixels, axis=0))
    k_eff = min(k, unique_count)
    if k_eff < 1:
        return []

    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 0.5)
    _, labels, centers = cv2.kmeans(pixels.astype(np.float32), k_eff, None,
                                     criteria, 3, cv2.KMEANS_PP_CENTERS)

    labels = labels.flatten()
    counts = np.bincount(labels, minlength=k_eff)
    total = counts.sum()

    colors = []
    for i in range(k_eff):
        b, g, r = centers[i]
        colors.append({
            "rgb": [int(round(r)), int(round(g)), int(round(b))],
            "pct": round(float(counts[i]) / total * 100, 1),
        })
    colors.sort(key=lambda c: c["pct"], reverse=True)
    return colors


def _lab_distance(rgb_a, rgb_b):
    a = np.uint8([[[rgb_a[2], rgb_a[1], rgb_a[0]]]])  # rgb -> bgr for cv2
    b = np.uint8([[[rgb_b[2], rgb_b[1], rgb_b[0]]]])
    lab_a = cv2.cvtColor(a, cv2.COLOR_BGR2LAB)[0][0].astype(np.float32)
    lab_b = cv2.cvtColor(b, cv2.COLOR_BGR2LAB)[0][0].astype(np.float32)
    return float(np.linalg.norm(lab_a - lab_b))


def match_color_pairs(input_colors, template_colors):
    """For each of the input's dominant colors, find the closest template
    dominant color by perceptual (Lab) distance and report a similarity %."""
    pairs = []
    for ic in input_colors:
        if not template_colors:
            break
        best = min(template_colors, key=lambda tc: _lab_distance(ic["rgb"], tc["rgb"]))
        dist = _lab_distance(ic["rgb"], best["rgb"])
        similarity = max(0.0, 100.0 * (1 - dist / MAX_LAB_DISTANCE))
        pairs.append({
            "input_rgb": ic["rgb"],
            "input_pct": ic["pct"],
            "template_rgb": best["rgb"],
            "template_pct": best["pct"],
            "similarity": round(similarity, 1),
        })
    return pairs


def set_template_color_data(name, bgr, mask):
    _template_color_data[name] = {
        "hist": compute_hs_histogram(bgr, mask),
        "colors": dominant_colors(bgr, mask),
    }


def compare_input_to_templates(input_bgr, input_mask):
    input_hist = compute_hs_histogram(input_bgr, input_mask)
    input_colors = dominant_colors(input_bgr, input_mask)

    results = []
    for name, data in _template_color_data.items():
        match_pct = histogram_match_pct(input_hist, data["hist"])
        pairs = match_color_pairs(input_colors, data["colors"])
        results.append({
            "template": name,
            "match_pct": match_pct,
            "color_pairs": pairs,
        })
    results.sort(key=lambda r: r["match_pct"], reverse=True)

    return {
        "input_dominant_colors": input_colors,
        "templates": results,
    }