#!/usr/bin/env python3
"""
colormatch.py - Farbabgleich zwischen einem Referenzbild (z. B. UARTScope)
und einem Testbild (z. B. HDMIViewer), das zu dunkel/anders erscheint.

Es wird KEINE Pixel-fuer-Pixel-Gleichheit verlangt (die Bilder duerfen
unterschiedliche Aufloesung und leicht unterschiedlichen Inhalt haben, z. B.
eine andere Animationsphase). Stattdessen werden pro Farbkanal die
Helligkeitsverteilungen (Quantile) beider Bilder zugeordnet und daraus die
Korrektur bestimmt, die das Testbild an die Referenz angleicht:

    ref  ~=  T(test)

Gefittet werden drei Modelle pro Kanal (Werte normiert auf 0..1):
    gain   :  out = g * in
    linear :  out = g * in + o
    gamma  :  out = g * in ** gamma

Ausgabe: Werte, Fehler (RMS in 0..255) und fertige GLSL-Zeilen fuer
VideoWidget.cpp. Optional LUT als CSV und das korrigierte Testbild.

WICHTIG: Fuer eine sinnvolle Aussage muessen beide Bilder auf demselben Weg
entstanden sein, z. B. beide per Desktop-Screenshot (Spectacle o. ae.) des
LAUFENDEN Programms, jeweils auf das Videobild zugeschnitten. Ein per P
gespeicherter HDMIViewer-Screenshot enthaelt nur den Framebuffer und zeigt
Unterschiede, die erst bei der Anzeige am Bildschirm entstehen, nicht.

Ist die Verarbeitung nicht linear (Kontrast, Studio-Range ...), liefert der
Fit nur eine Naeherung (Fehler deutlich ueber ~3). Dann hilft --palette: Die
haeufigsten Flaechenfarben zeigen direkt, was mit Weiss/Rot/... passiert.

Aufruf:
    colormatch.py REFERENZ.png TEST.png
    colormatch.py ref.png test.png --palette
    colormatch.py ref.png test.png --crop-ref 100,80,1280,720 --crop-test 0,0,1280,672
    colormatch.py ref.png test.png --lut lut.csv --save-corrected korrigiert.png

Benoetigt: python3, numpy, Pillow  (Arch: pacman -S python-numpy python-pillow)
"""

import argparse
import sys

import numpy as np
from PIL import Image

CHANNELS = ("R", "G", "B")
N_QUANT = 4096          # Anzahl der Quantil-Stuetzstellen pro Kanal


def parse_crop(text):
    if not text:
        return None
    try:
        x, y, w, h = (int(v) for v in text.split(","))
    except ValueError:
        raise argparse.ArgumentTypeError("Format: x,y,breite,hoehe")
    return x, y, w, h


def load(path, crop, auto, thr):
    arr = np.asarray(Image.open(path).convert("RGB"))
    if crop:
        x, y, w, h = crop
        arr = arr[y:y + h, x:x + w]
    if auto:
        # Schwarze Letterbox-/Pillarbox-Raender abschneiden
        mask = arr.max(axis=2) > thr
        ys = np.where(mask.any(axis=1))[0]
        xs = np.where(mask.any(axis=0))[0]
        if len(ys) and len(xs):
            arr = arr[ys[0]:ys[-1] + 1, xs[0]:xs[-1] + 1]
    if arr.size == 0:
        sys.exit(f"Leerer Ausschnitt bei {path}")
    return arr


def quantile_pairs(test_ch, ref_ch):
    """Ordnet Quantile des Testkanals denen des Referenzkanals zu.
    Rueckgabe: (t, r, w) - eindeutige Wertepaare (0..1) mit Gewicht."""
    ts = np.sort(test_ch.ravel())
    rs = np.sort(ref_ch.ravel())
    p = (np.arange(N_QUANT) + 0.5) / N_QUANT
    t = ts[(p * (len(ts) - 1)).astype(np.int64)].astype(np.float64) / 255.0
    r = rs[(p * (len(rs) - 1)).astype(np.int64)].astype(np.float64) / 255.0
    pairs = np.stack([t, r], axis=1)
    uniq, cnt = np.unique(pairs, axis=0, return_counts=True)
    return uniq[:, 0], uniq[:, 1], cnt.astype(np.float64)


def fit_gain(t, r, w):
    den = np.sum(w * t * t)
    g = np.sum(w * t * r) / den if den > 0 else 1.0
    return {"g": g}


def fit_linear(t, r, w):
    sw = np.sqrt(w)
    a = np.stack([t, np.ones_like(t)], axis=1) * sw[:, None]
    sol, *_ = np.linalg.lstsq(a, r * sw, rcond=None)
    return {"g": sol[0], "o": sol[1]}


def fit_gamma(t, r, w):
    m = (t > 1.0 / 255.0) & (r > 1.0 / 255.0)
    if m.sum() < 2:
        return {"g": 1.0, "gamma": 1.0}
    sw = np.sqrt(w[m])
    a = np.stack([np.ones(m.sum()), np.log(t[m])], axis=1) * sw[:, None]
    sol, *_ = np.linalg.lstsq(a, np.log(r[m]) * sw, rcond=None)
    return {"g": float(np.exp(sol[0])), "gamma": float(sol[1])}


def apply_model(name, p, x):
    if name == "gain":
        y = p["g"] * x
    elif name == "linear":
        y = p["g"] * x + p["o"]
    else:
        y = p["g"] * np.power(np.clip(x, 0.0, 1.0), p["gamma"])
    return np.clip(y, 0.0, 1.0)


def rms255(name, p, t, r, w):
    d = (apply_model(name, p, t) - r) * 255.0
    return float(np.sqrt(np.sum(w * d * d) / np.sum(w)))


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("ref", help="Referenzbild (so soll es aussehen, z. B. UARTScope)")
    ap.add_argument("test", help="Testbild (so sieht es aus, z. B. HDMIViewer)")
    ap.add_argument("--crop-ref", type=parse_crop, help="Ausschnitt x,y,b,h im Referenzbild")
    ap.add_argument("--crop-test", type=parse_crop, help="Ausschnitt x,y,b,h im Testbild")
    ap.add_argument("--no-autocrop", action="store_true",
                    help="schwarze Raender NICHT automatisch abschneiden")
    ap.add_argument("--black-threshold", type=int, default=8,
                    help="Schwellwert fuer die Randerkennung (Standard 8)")
    ap.add_argument("--palette", action="store_true",
                    help="zusaetzlich die haeufigsten Flaechenfarben beider Bilder ausgeben")
    ap.add_argument("--lut", metavar="CSV", help="256-Werte-LUT pro Kanal als CSV speichern")
    ap.add_argument("--save-corrected", metavar="PNG",
                    help="Testbild mit der besten Korrektur speichern")
    args = ap.parse_args()

    auto = not args.no_autocrop
    ref = load(args.ref, args.crop_ref, auto, args.black_threshold)
    test = load(args.test, args.crop_test, auto, args.black_threshold)

    print(f"Referenz: {ref.shape[1]}x{ref.shape[0]}   Test: {test.shape[1]}x{test.shape[0]}")
    print()

    if args.palette:
        # Haeufigste Farben (ohne fast schwarze Pixel). Bei Bildern mit flachen
        # Farbflaechen (Amiga, Konsolen) zeigt das am direktesten, wie die
        # Helligkeit veraendert wird, z. B. Weiss 255 -> 128 = Faktor 0,5.
        for label, img in (("Referenz", ref), ("Test", test)):
            px = img.reshape(-1, 3)
            px = px[px.max(axis=1) > 8]
            u, c = np.unique(px, axis=0, return_counts=True)
            print(f"Haeufigste Farben {label}:")
            for i in np.argsort(-c)[:6]:
                print(f"    ({u[i][0]:3d},{u[i][1]:3d},{u[i][2]:3d})  {c[i]} Pixel")
        print()

    print("Mittelwerte (0..255)        R        G        B")
    mr = ref.reshape(-1, 3).mean(axis=0)
    mt = test.reshape(-1, 3).mean(axis=0)
    print("  Referenz            " + "".join(f"{v:9.2f}" for v in mr))
    print("  Test                " + "".join(f"{v:9.2f}" for v in mt))
    print("  Verhaeltnis Ref/Test" + "".join(f"{(a / b if b > 0 else float('nan')):9.3f}"
                                              for a, b in zip(mr, mt)))
    print()

    results = {}
    luts = []
    for ci, cname in enumerate(CHANNELS):
        t, r, w = quantile_pairs(test[..., ci], ref[..., ci])
        models = {
            "gain": fit_gain(t, r, w),
            "linear": fit_linear(t, r, w),
            "gamma": fit_gamma(t, r, w),
        }
        errs = {n: rms255(n, p, t, r, w) for n, p in models.items()}
        results[cname] = (models, errs)

        # LUT (monoton, per Interpolation der Quantilpaare)
        order = np.argsort(t)
        tt, rr = t[order], r[order]
        xs = np.arange(256) / 255.0
        lut = np.interp(xs, tt, rr)
        lut = np.maximum.accumulate(lut)
        luts.append(np.clip(np.round(lut * 255.0), 0, 255).astype(np.uint8))

    print("Fit pro Kanal (Fehler = RMS in 0..255, kleiner ist besser)")
    for cname in CHANNELS:
        models, errs = results[cname]
        print(f"  {cname}:")
        print(f"    gain   : g={models['gain']['g']:.4f}"
              f"                         RMS {errs['gain']:.2f}")
        print(f"    linear : g={models['linear']['g']:.4f}  o={models['linear']['o']:+.4f}"
              f"            RMS {errs['linear']:.2f}")
        print(f"    gamma  : g={models['gamma']['g']:.4f}  gamma={models['gamma']['gamma']:.4f}"
              f"      RMS {errs['gamma']:.2f}")

    # Bestes gemeinsames Modell nach Summe der Fehlerquadrate waehlen,
    # aber das einfachste nehmen, wenn es fast genauso gut ist.
    total = {n: np.sqrt(np.mean([results[c][1][n] ** 2 for c in CHANNELS]))
             for n in ("gain", "linear", "gamma")}
    best = min(total, key=total.get)
    if total["gain"] <= total[best] + 1.0:
        best = "gain"
    elif total["linear"] <= total[best] + 1.0 and best == "gamma":
        best = "linear"

    print()
    print(f"Gesamtfehler  gain {total['gain']:.2f}   linear {total['linear']:.2f}"
          f"   gamma {total['gamma']:.2f}   ->  empfohlen: {best}")
    if total[best] > 4.0:
        print("ACHTUNG: Fehler ist hoch - die Abweichung ist wohl nicht durch ein einfaches")
        print("Modell erklaerbar (oder die Bildinhalte unterscheiden sich zu stark).")
        print("Mit --palette die Flaechenfarben ansehen, bevor die Werte uebernommen werden.")
    print()
    print("GLSL-Zeile fuer VideoWidget.cpp (im Fragment-Shader auf 'color' anwenden):")
    P = [results[c][0][best] for c in CHANNELS]
    if best == "gain":
        v = ", ".join(f"{p['g']:.4f}" for p in P)
        print(f"    color = clamp(color * vec3({v}), 0.0, 1.0);")
    elif best == "linear":
        g = ", ".join(f"{p['g']:.4f}" for p in P)
        o = ", ".join(f"{p['o']:.4f}" for p in P)
        print(f"    color = clamp(color * vec3({g}) + vec3({o}), 0.0, 1.0);")
    else:
        g = ", ".join(f"{p['g']:.4f}" for p in P)
        e = ", ".join(f"{p['gamma']:.4f}" for p in P)
        print(f"    color = clamp(vec3({g}) * pow(color, vec3({e})), 0.0, 1.0);")

    if args.lut:
        with open(args.lut, "w", encoding="utf-8") as f:
            f.write("in,R,G,B\n")
            for i in range(256):
                f.write(f"{i},{luts[0][i]},{luts[1][i]},{luts[2][i]}\n")
        print(f"\nLUT gespeichert: {args.lut}")

    if args.save_corrected:
        x = test.astype(np.float64) / 255.0
        out = np.empty_like(x)
        for ci, cname in enumerate(CHANNELS):
            out[..., ci] = apply_model(best, results[cname][0][best], x[..., ci])
        Image.fromarray(np.round(out * 255.0).astype(np.uint8)).save(args.save_corrected)
        print(f"Korrigiertes Testbild gespeichert: {args.save_corrected}")

    print()
    print("Hinweis: Weichen die Mittelwerte stark ab, obwohl beide Bilder denselben")
    print("Inhalt zeigen, ist die Korrektur sinnvoll. Sind die Bilder praktisch")
    print("gleich (Verhaeltnis ~1.0, Gain ~1.0), liegt der Unterschied NICHT in")
    print("den Bilddaten - dann haben beide Bilder denselben Weg genommen.")


if __name__ == "__main__":
    main()