#!/usr/bin/env python3
"""Reproduce the numerical tables in the transfer-function article.

Requirements: Python 3.10+ and NumPy.
All signal values are normalized to [0, 1] unless a luminance unit is stated.
The script performs no image I/O and does not depend on a colour-management library.
"""
from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable

import numpy as np


def srgb_decode(v: np.ndarray | float) -> np.ndarray:
    """IEC 61966-2-1 encoded sRGB -> linear relative RGB."""
    x = np.asarray(v, dtype=np.float64)
    return np.where(x <= 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4)


def srgb_encode(l: np.ndarray | float) -> np.ndarray:
    """Linear relative RGB -> IEC 61966-2-1 encoded sRGB."""
    x = np.asarray(l, dtype=np.float64)
    return np.where(x <= 0.0031308, 12.92 * x, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)


def bt709_oetf(e: np.ndarray | float) -> np.ndarray:
    x = np.asarray(e, dtype=np.float64)
    return np.where(x < 0.018, 4.5 * x, 1.099 * np.power(x, 0.45) - 0.099)


def bt709_inverse_oetf(ep: np.ndarray | float) -> np.ndarray:
    x = np.asarray(ep, dtype=np.float64)
    return np.where(x < 0.081, x / 4.5, np.power((x + 0.099) / 1.099, 1.0 / 0.45))


def bt1886_parameters(lw: float, lb: float, gamma: float = 2.4) -> tuple[float, float]:
    if not (lw > lb >= 0):
        raise ValueError("Require Lw > Lb >= 0")
    lwg = lw ** (1.0 / gamma)
    lbg = lb ** (1.0 / gamma)
    a = (lwg - lbg) ** gamma
    b = lbg / (lwg - lbg)
    return a, b


def bt1886_eotf(v: np.ndarray | float, lw: float = 100.0, lb: float = 0.0,
                 gamma: float = 2.4) -> np.ndarray:
    x = np.asarray(v, dtype=np.float64)
    a, b = bt1886_parameters(lw, lb, gamma)
    return a * np.maximum(x + b, 0.0) ** gamma


def bt1886_inverse(l: np.ndarray | float, lw: float = 100.0, lb: float = 0.0,
                   gamma: float = 2.4) -> np.ndarray:
    x = np.asarray(l, dtype=np.float64)
    a, b = bt1886_parameters(lw, lb, gamma)
    return np.maximum(np.power(np.maximum(x, 0.0) / a, 1.0 / gamma) - b, 0.0)


def pq_eotf(ep: np.ndarray | float) -> np.ndarray:
    """BT.2100/ST 2084 PQ code -> cd/m^2."""
    x = np.asarray(ep, dtype=np.float64)
    m1 = 2610.0 / 16384.0
    m2 = (2523.0 / 4096.0) * 128.0
    c1 = 3424.0 / 4096.0
    c2 = (2413.0 / 4096.0) * 32.0
    c3 = (2392.0 / 4096.0) * 32.0
    p = np.power(np.maximum(x, 0.0), 1.0 / m2)
    y = np.power(np.maximum(p - c1, 0.0) / (c2 - c3 * p), 1.0 / m1)
    return 10000.0 * y


def pq_inverse_eotf(luminance_cd_m2: np.ndarray | float) -> np.ndarray:
    """cd/m^2 -> BT.2100/ST 2084 PQ code."""
    l = np.asarray(luminance_cd_m2, dtype=np.float64) / 10000.0
    m1 = 2610.0 / 16384.0
    m2 = (2523.0 / 4096.0) * 128.0
    c1 = 3424.0 / 4096.0
    c2 = (2413.0 / 4096.0) * 32.0
    c3 = (2392.0 / 4096.0) * 32.0
    lm = np.power(np.maximum(l, 0.0), m1)
    return np.power((c1 + c2 * lm) / (1.0 + c3 * lm), m2)


def hlg_oetf(e: np.ndarray | float) -> np.ndarray:
    x = np.asarray(e, dtype=np.float64)
    a = 0.17883277
    b = 0.28466892
    c = 0.55991073
    return np.where(x <= 1.0 / 12.0, np.sqrt(3.0 * x), a * np.log(12.0 * x - b) + c)


def hlg_inverse_oetf(ep: np.ndarray | float) -> np.ndarray:
    x = np.asarray(ep, dtype=np.float64)
    a = 0.17883277
    b = 0.28466892
    c = 0.55991073
    return np.where(x <= 0.5, x * x / 3.0, (np.exp((x - c) / a) + b) / 12.0)


def hlg_system_gamma(lw_cd_m2: float) -> float:
    return 1.2 + 0.42 * math.log10(lw_cd_m2 / 1000.0)


M709_TO_XYZ_D65 = np.array(
    [[0.4123907992659595, 0.3575843393838780, 0.1804807884018343],
     [0.2126390058715104, 0.7151686787677559, 0.0721923153607337],
     [0.0193308187155919, 0.1191947797946260, 0.9505321522496607]],
    dtype=np.float64,
)
D65_XYZ = np.array([0.9504559270516716, 1.0, 1.0890577507598784])


def linear_rgb_to_lab(rgb: np.ndarray) -> np.ndarray:
    rgb = np.asarray(rgb, dtype=np.float64)
    xyz = rgb @ M709_TO_XYZ_D65.T
    r = xyz / D65_XYZ
    delta = 6.0 / 29.0
    f = np.where(r > delta**3, np.cbrt(r), r / (3.0 * delta**2) + 4.0 / 29.0)
    return np.stack(
        [116.0 * f[..., 1] - 16.0,
         500.0 * (f[..., 0] - f[..., 1]),
         200.0 * (f[..., 1] - f[..., 2])],
        axis=-1,
    )


def delta_e_2000(lab1: np.ndarray, lab2: np.ndarray) -> np.ndarray:
    """Vectorized CIEDE2000, kL=kC=kH=1."""
    l1, a1, b1 = np.moveaxis(np.asarray(lab1, dtype=np.float64), -1, 0)
    l2, a2, b2 = np.moveaxis(np.asarray(lab2, dtype=np.float64), -1, 0)
    c1 = np.hypot(a1, b1)
    c2 = np.hypot(a2, b2)
    cbar = (c1 + c2) / 2.0
    g = 0.5 * (1.0 - np.sqrt(cbar**7 / (cbar**7 + 25.0**7)))
    a1p = (1.0 + g) * a1
    a2p = (1.0 + g) * a2
    c1p = np.hypot(a1p, b1)
    c2p = np.hypot(a2p, b2)
    h1p = np.mod(np.degrees(np.arctan2(b1, a1p)), 360.0)
    h2p = np.mod(np.degrees(np.arctan2(b2, a2p)), 360.0)

    dl = l2 - l1
    dc = c2p - c1p
    dh_raw = h2p - h1p
    both_nonzero = (c1p * c2p) != 0
    dh_deg = np.where(
        ~both_nonzero, 0.0,
        np.where(np.abs(dh_raw) <= 180.0, dh_raw,
                 np.where(dh_raw > 180.0, dh_raw - 360.0, dh_raw + 360.0)),
    )
    dh = 2.0 * np.sqrt(c1p * c2p) * np.sin(np.radians(dh_deg) / 2.0)

    lbar = (l1 + l2) / 2.0
    cbarp = (c1p + c2p) / 2.0
    hsum = h1p + h2p
    hbar = np.where(
        ~both_nonzero, hsum,
        np.where(np.abs(h1p - h2p) <= 180.0, hsum / 2.0,
                 np.where(hsum < 360.0, (hsum + 360.0) / 2.0, (hsum - 360.0) / 2.0)),
    )
    t = (1.0
         - 0.17 * np.cos(np.radians(hbar - 30.0))
         + 0.24 * np.cos(np.radians(2.0 * hbar))
         + 0.32 * np.cos(np.radians(3.0 * hbar + 6.0))
         - 0.20 * np.cos(np.radians(4.0 * hbar - 63.0)))
    sl = 1.0 + 0.015 * (lbar - 50.0) ** 2 / np.sqrt(20.0 + (lbar - 50.0) ** 2)
    sc = 1.0 + 0.045 * cbarp
    sh = 1.0 + 0.015 * cbarp * t
    dtheta = 30.0 * np.exp(-((hbar - 275.0) / 25.0) ** 2)
    rc = 2.0 * np.sqrt(cbarp**7 / (cbarp**7 + 25.0**7))
    rt = -rc * np.sin(np.radians(2.0 * dtheta))
    return np.sqrt((dl / sl) ** 2 + (dc / sc) ** 2 + (dh / sh) ** 2
                   + rt * (dc / sc) * (dh / sh))


def rgb_to_ycbcr(rgb: Iterable[float], kr: float, kb: float) -> np.ndarray:
    r, g, b = np.asarray(tuple(rgb), dtype=np.float64)
    kg = 1.0 - kr - kb
    y = kr * r + kg * g + kb * b
    cb = (b - y) / (2.0 * (1.0 - kb))
    cr = (r - y) / (2.0 * (1.0 - kr))
    return np.array([y, cb, cr])


def ycbcr_to_rgb(ycc: Iterable[float], kr: float, kb: float) -> np.ndarray:
    y, cb, cr = np.asarray(tuple(ycc), dtype=np.float64)
    kg = 1.0 - kr - kb
    r = y + 2.0 * (1.0 - kr) * cr
    b = y + 2.0 * (1.0 - kb) * cb
    g = (y - kr * r - kb * b) / kg
    return np.array([r, g, b])


def levels_gamma(x: np.ndarray | float, gamma: float) -> np.ndarray:
    """Normalized AviSynth/VapourSynth Levels gamma core: y=x^(1/gamma)."""
    return np.power(np.clip(np.asarray(x, dtype=np.float64), 0.0, 1.0), 1.0 / gamma)


def table_srgb_vs_22() -> list[dict[str, float | int | None]]:
    codes = np.array([0, 1, 2, 4, 8, 12, 16, 20, 32, 64, 128, 192, 255])
    v = codes / 255.0
    s = srgb_decode(v)
    p = v**2.2
    return [
        {"code": int(c), "sRGB_linear": float(a), "power_2_2_linear": float(b),
         "absolute_difference": float(b - a),
         "relative_difference_percent": None if a == 0 else float((b / a - 1.0) * 100.0),
         "difference_cd_m2_at_100_white": float((b - a) * 100.0)}
        for c, a, b in zip(codes, s, p, strict=True)
    ]


def table_srgb_to_gamma24() -> tuple[list[dict[str, float | int]], dict[str, float | int]]:
    codes = np.array([0, 8, 16, 32, 64, 128, 192, 255])
    v = codes / 255.0
    exact = srgb_decode(v) ** (1.0 / 2.4)
    approx = v ** (2.2 / 2.4)
    rows = [
        {"input_code": int(c), "exact_target_code": float(e * 255.0),
         "power_ratio_target_code": float(a * 255.0),
         "signed_error_code": float((a - e) * 255.0)}
        for c, e, a in zip(codes, exact, approx, strict=True)
    ]
    all_v = np.arange(256, dtype=np.float64) / 255.0
    all_err = (all_v ** (2.2 / 2.4) - srgb_decode(all_v) ** (1.0 / 2.4)) * 255.0
    metrics = {
        "maximum_absolute_code_error": float(np.max(np.abs(all_err))),
        "input_code_at_maximum": int(np.argmax(np.abs(all_err))),
        "mean_absolute_code_error": float(np.mean(np.abs(all_err))),
        "rmse_code": float(np.sqrt(np.mean(all_err**2))),
        "mean_signed_code_error": float(np.mean(all_err)),
        "p95_absolute_code_error": float(np.quantile(np.abs(all_err), 0.95)),
    }
    return rows, metrics


def table_bt1886_black() -> list[dict[str, float | int | str]]:
    x = np.arange(256, dtype=np.float64) / 255.0
    approximation = x ** (2.2 / 2.4)
    rows: list[dict[str, float | int | str]] = []
    for lb in (0.0, 0.001, 0.01, 0.1, 1.0):
        lw = 100.0
        desired_l = lb + (lw - lb) * x**2.2
        exact_v = bt1886_inverse(desired_l, lw, lb)
        error = (approximation - exact_v) * 255.0
        rows.append({
            "Lw_cd_m2": lw,
            "Lb_cd_m2": lb,
            "contrast": "infinite" if lb == 0 else f"{lw/lb:.0f}:1",
            "maximum_absolute_code_error": float(np.max(np.abs(error))),
            "input_code_at_maximum": int(np.argmax(np.abs(error))),
            "mean_absolute_code_error": float(np.mean(np.abs(error))),
            "rmse_code": float(np.sqrt(np.mean(error**2))),
        })
    return rows


def perceptual_metrics() -> dict[str, object]:
    values = np.arange(256, dtype=np.float64) / 255.0
    exact_gray = np.repeat(srgb_decode(values)[:, None], 3, axis=1)
    approximate_gray = np.repeat((values**2.2)[:, None], 3, axis=1)
    de_gray = delta_e_2000(linear_rgb_to_lab(exact_gray), linear_rgb_to_lab(approximate_gray))

    grid = np.linspace(0.0, 1.0, 17)
    code_rgb = np.array(np.meshgrid(grid, grid, grid, indexing="ij")).reshape(3, -1).T
    de_cube = delta_e_2000(
        linear_rgb_to_lab(srgb_decode(code_rgb)),
        linear_rgb_to_lab(code_rgb**2.2),
    )
    i = int(np.argmax(de_cube))
    return {
        "neutral_256": {
            "mean_deltaE00": float(np.mean(de_gray)),
            "p95_deltaE00": float(np.quantile(de_gray, 0.95)),
            "maximum_deltaE00": float(np.max(de_gray)),
            "code_at_maximum": int(np.argmax(de_gray)),
            "fraction_deltaE00_ge_1": float(np.mean(de_gray >= 1.0)),
        },
        "rgb_cube_17": {
            "sample_count": int(len(de_cube)),
            "mean_deltaE00": float(np.mean(de_cube)),
            "p95_deltaE00": float(np.quantile(de_cube, 0.95)),
            "p99_deltaE00": float(np.quantile(de_cube, 0.99)),
            "maximum_deltaE00": float(np.max(de_cube)),
            "code_rgb_at_maximum": [float(v) for v in code_rgb[i]],
            "fraction_deltaE00_ge_1": float(np.mean(de_cube >= 1.0)),
        },
    }


def y_only_vs_rgb_example() -> dict[str, object]:
    rgb = np.array([0.8, 0.4, 0.2])
    kr, kb = 0.2126, 0.0722
    ycc = rgb_to_ycbcr(rgb, kr, kb)
    ycc_y = ycc.copy()
    ycc_y[0] = levels_gamma(ycc[0], 0.88)
    rgb_y_only = ycbcr_to_rgb(ycc_y, kr, kb)
    rgb_per_channel = levels_gamma(rgb, 0.88)
    return {
        "input_RpGpBp": [float(v) for v in rgb],
        "input_YpCbCr": [float(v) for v in ycc],
        "gamma_parameter": 0.88,
        "mathematical_exponent": 1.0 / 0.88,
        "Y_only_result_RpGpBp": [float(v) for v in rgb_y_only],
        "RGB_per_channel_result_RpGpBp": [float(v) for v in rgb_per_channel],
    }


def system_gamma_samples() -> list[dict[str, float]]:
    e = np.array([0.018, 0.02, 0.05, 0.10, 0.18, 0.25, 0.50, 0.75, 1.0])
    ep = bt709_oetf(e)
    f = ep**2.4
    local = np.where(
        e < 0.018,
        2.4,
        2.4 * (1.099 * 0.45 * e**0.45) / ep,
    )
    return [
        {"scene_linear_E": float(a), "BT709_code_Eprime": float(b),
         "BT1886_output_relative": float(c), "local_log_slope": float(d)}
        for a, b, c, d in zip(e, ep, f, local, strict=True)
    ]


def main() -> None:
    srgb_to_24_rows, srgb_to_24_metrics = table_srgb_to_gamma24()
    payload = {
        "assumptions": {
            "normalization": "Signals normalized to [0,1] unless stated",
            "white_for_srgb_table_cd_m2": 100,
            "CIEDE2000_white": "D65",
            "CIEDE2000_primaries": "BT.709/sRGB",
            "quantization_after_conversion": False,
        },
        "srgb_vs_pure_2_2": table_srgb_vs_22(),
        "srgb_to_pure_gamma_2_4": {
            "selected_codes": srgb_to_24_rows,
            "all_256_codes_metrics": srgb_to_24_metrics,
        },
        "bt1886_nonzero_black": table_bt1886_black(),
        "perceptual_metrics": perceptual_metrics(),
        "y_only_vs_rgb_gamma": y_only_vs_rgb_example(),
        "bt709_plus_bt1886_local_system_gamma": system_gamma_samples(),
        "checkerboard": {
            "gamma_space_average_code": 0.5,
            "sRGB_linear_luminance_of_code_0_5": float(srgb_decode(0.5)),
            "correct_code_for_50_percent_linear_light": float(srgb_encode(0.5)),
        },
        "hlg_system_gamma_examples": [
            {"Lw_cd_m2": lw, "gamma": hlg_system_gamma(lw)}
            for lw in (400, 600, 800, 1000, 1500, 2000)
        ],
    }
    out = Path(__file__).with_name("gamma-analysis-results.json")
    out.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(json.dumps(payload, indent=2, ensure_ascii=False))
    print(f"\nWrote {out}")


if __name__ == "__main__":
    main()