#!/usr/bin/env python3
"""Generate the skill's .cube LUTs: a D-Log(M) conversion + Rec.709 looks.

TWO FAMILIES — never confuse them (a conversion LUT on non-log footage
crushes blacks, blows highlights, and pushes colors out of gamut):

  conversion/dji-dlogm-to-rec709.cube   input MUST be D-Log/D-Log M footage.
                                        Applied automatically by
                                        normalize_clips.py to dlog clips only.
  looks/reel-{gritty,punchy,faded}.cube Rec.709 in / Rec.709 out — safe on any
                                        normalized or SDR footage, in ffmpeg
                                        (render_reel.py --lut) or an NLE node.

(Resolve's LUT-browser thumbnails are Rec.709 stills, so the conversion LUT
always previews as blown-out neon there — expected, not a defect. Same for
DJI's official LUT.)

Conversion math is DJI's official "White Paper on D-Log and D-Gamut"
(dl.djicdn.com, rev 1.0 2017): exact D-Log-to-linear transfer + the
D-Gamut-to-Rec.709 matrix, then a filmic tonemap (ACES fit) and 2.2 display
encode. DJI has never published the D-Log M curve (only per-model conversion
LUTs on dji.com), so for D-Log M this is an APPROXIMATION — for exactness,
download the official "D-Log M to Rec.709" .cube for the user's drone and use
it directly (normalize_clips.py --dlog-lut official.cube), or bake a look on
top of it:

  make_luts.py --compose official_dlogm.cube --looks gritty --outdir work/luts

which samples the official LUT as the base and applies our look after it —
accurate conversion, branded look, no redistribution of DJI's file. The same
one-node combined output (whitepaper base) is available via --combined.

Stdlib only — no numpy/colour-science needed at this cube size.

Usage:
  make_luts.py                                  # regenerate the bundled set
  make_luts.py --looks gritty,punchy --size 25 --outdir work/luts
  make_luts.py --compose official.cube --looks all --outdir work/luts
"""
import argparse, math, sys
from pathlib import Path

# --- DJI whitepaper math (D-Log / D-Gamut, verbatim constants) ---

def dlog_to_linear(x):
    if x <= 0.14:
        return (x - 0.0929) / 6.025
    return (10 ** (3.89616 * x - 2.27752) - 0.0108) / 0.9892

DGAMUT_TO_R709 = ((1.6746, -0.5797, -0.0949),
                  (-0.0981, 1.3340, -0.2359),
                  (-0.0410, -0.2430, 1.2840))

def mat3(m, rgb):
    r, g, b = rgb
    return tuple(m[i][0] * r + m[i][1] * g + m[i][2] * b for i in range(3))

# --- display rendering ---

EXPOSURE = 0.7   # anchors D-Log 18% gray (code 0.3988) at ~0.45 display

def aces_fit(x):  # Narkowicz ACES filmic approximation, soft shoulder
    x = max(0.0, x) * EXPOSURE
    return (x * (2.51 * x + 0.03)) / (x * (2.43 * x + 0.59) + 0.14)

def encode(x):
    return min(1.0, max(0.0, x)) ** (1 / 2.2)

def neutral(rgb):
    lin = tuple(dlog_to_linear(c) for c in rgb)
    lin = mat3(DGAMUT_TO_R709, lin)
    return tuple(encode(aces_fit(c)) for c in lin)

# --- style ops (display space) ---

def luma(rgb):
    return 0.2126 * rgb[0] + 0.7152 * rgb[1] + 0.0722 * rgb[2]

def saturate(rgb, s):
    y = luma(rgb)
    return tuple(y + (c - y) * s for c in rgb)

def contrast(rgb, amt):  # blend toward smoothstep S-curve
    return tuple(c + (c * c * (3 - 2 * c) - c) * amt for c in rgb)

def lift_gain(rgb, lift, gain):  # per-channel: teal shadows, warm mids...
    return tuple(g * (c + l * (1 - c)) for c, l, g in zip(rgb, lift, gain))

def clamp(rgb):
    return tuple(min(1.0, max(0.0, c)) for c in rgb)

STYLES = {
    # neutral: technical conversion only
    "neutral": lambda rgb: rgb,
    # gritty: the skill's teal/orange "Gritty Authentic Energy" as a LUT
    "gritty": lambda rgb: clamp(contrast(saturate(
        lift_gain(rgb, (-0.010, 0.000, 0.028), (1.035, 1.000, 0.955)), 1.12), 0.22)),
    # punchy: strong S-curve + saturation, no hue shift
    "punchy": lambda rgb: clamp(contrast(saturate(rgb, 1.18), 0.42)),
    # faded: lifted matte blacks, muted color, warm highlights
    "faded": lambda rgb: clamp(saturate(
        lift_gain(rgb, (0.055, 0.050, 0.045), (1.015, 1.000, 0.960)), 0.88)),
}

# --- optional base: sample a downloaded official .cube instead of our math ---

def read_cube(path):
    size, table = None, []
    for line in Path(path).read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or line.startswith(("TITLE", "DOMAIN")):
            continue
        if line.startswith("LUT_3D_SIZE"):
            size = int(line.split()[-1])
            continue
        if line.startswith("LUT_1D_SIZE"):
            sys.exit(f"{path} is a 1D LUT — need a 3D .cube")
        parts = line.split()
        if len(parts) == 3:
            table.append(tuple(float(v) for v in parts))
    if not size or len(table) != size ** 3:
        sys.exit(f"Malformed cube {path}: size={size}, rows={len(table)}")
    return size, table

def cube_sampler(path):
    size, table = read_cube(path)
    n = size - 1
    def at(ri, gi, bi):
        return table[bi * size * size + gi * size + ri]  # red fastest
    def sample(rgb):  # trilinear
        out = [0.0, 0.0, 0.0]
        pos = [min(max(c, 0.0), 1.0) * n for c in rgb]
        i0 = [min(int(p), n - 1) if n else 0 for p in pos]
        f = [p - i for p, i in zip(pos, i0)]
        for dr in (0, 1):
            for dg in (0, 1):
                for db in (0, 1):
                    w = ((f[0] if dr else 1 - f[0]) *
                         (f[1] if dg else 1 - f[1]) *
                         (f[2] if db else 1 - f[2]))
                    v = at(i0[0] + dr, i0[1] + dg, i0[2] + db)
                    for k in range(3):
                        out[k] += w * v[k]
        return tuple(out)
    return sample


def identity(rgb):
    return rgb


def write_cube(path, size, transform, title):
    n = size - 1
    tmp = path.with_name(path.name + ".tmp")
    with tmp.open("w") as f:
        f.write(f'TITLE "{title}"\n')
        f.write(f"LUT_3D_SIZE {size}\n")
        for bi in range(size):              # red varies fastest (.cube spec)
            for gi in range(size):
                for ri in range(size):
                    rgb = transform((ri / n, gi / n, bi / n))
                    f.write(f"{rgb[0]:.6f} {rgb[1]:.6f} {rgb[2]:.6f}\n")
    tmp.rename(path)
    print(f"  {path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--outdir",
                    default=str(Path(__file__).resolve().parent.parent / "assets/luts"))
    ap.add_argument("--size", type=int, default=33)
    ap.add_argument("--looks", default="all",
                    help=f"Comma list of {','.join(k for k in STYLES if k != 'neutral')} "
                         "(default all)")
    ap.add_argument("--compose", metavar="OFFICIAL.cube",
                    help="Emit combined conversion+look LUTs using this (e.g. DJI's "
                         "official D-Log M) LUT as the base conversion")
    ap.add_argument("--combined", action="store_true",
                    help="Emit combined conversion+look LUTs from the whitepaper "
                         "base (one-node workflows; NOT for non-log footage)")
    args = ap.parse_args()

    looks = [k for k in STYLES if k != "neutral"] if args.looks == "all" \
        else args.looks.split(",")
    bad = [s for s in looks if s not in STYLES or s == "neutral"]
    if bad:
        sys.exit(f"Unknown look(s): {bad} — choose from "
                 f"{[k for k in STYLES if k != 'neutral']}")

    outdir = Path(args.outdir)
    size = args.size

    if args.compose or args.combined:
        # combined: conversion + look in one cube. Input must be D-Log(M).
        base = cube_sampler(args.compose) if args.compose else neutral
        tag = "official-base" if args.compose else "wp-approx"
        outdir.mkdir(parents=True, exist_ok=True)
        for style in looks:
            op = STYLES[style]
            write_cube(outdir / f"dji-dlogm-to-rec709-{style}.cube", size,
                       lambda rgb, op=op: op(base(rgb)),
                       f"DJI D-Log(M) -> Rec.709 + {style} look "
                       f"[reel-builder {tag}] — APPLY TO LOG FOOTAGE ONLY")
    else:
        # bundled set: one conversion LUT + Rec.709-safe looks, kept separate
        (outdir / "conversion").mkdir(parents=True, exist_ok=True)
        (outdir / "looks").mkdir(parents=True, exist_ok=True)
        write_cube(outdir / "conversion/dji-dlogm-to-rec709.cube", size, neutral,
                   "DJI D-Log(M) -> Rec.709 [reel-builder wp-approx] — "
                   "APPLY TO LOG FOOTAGE ONLY")
        for style in looks:
            write_cube(outdir / f"looks/reel-{style}.cube", size, STYLES[style],
                       f"Reel {style} look — Rec.709 in/out, any footage")

    # sanity: gray stays gray, tone is monotonic, mid-gray anchored, looks
    # are identity-anchored (black stays black-ish, white stays white-ish)
    g1, g2 = neutral((0.3988,) * 3), neutral((0.55,) * 3)
    assert max(g1) - min(g1) < 1e-6, "gray drifted — matrix rows must sum to 1"
    assert g2 > g1, "tone curve not monotonic"
    assert 0.40 < g1[0] < 0.50, f"mid-gray anchor off: {g1[0]:.3f}"
    for style in ("gritty", "punchy"):
        mid = STYLES[style]((0.5, 0.5, 0.5))
        assert all(0.35 < c < 0.65 for c in mid), f"{style} mid-tone off: {mid}"
    print(f"Sanity: D-Log 18% gray -> {g1[0]:.3f} display, monotonic, "
          "neutral-gray OK, looks mid-anchored OK")


if __name__ == "__main__":
    main()