"""GPU develop stack — port of probe_parity/develop.py render_srgb_develop."""

from __future__ import annotations

import torch
import torch.nn.functional as F

from ..polarr_color_space import ENCODING_PRESETS
from ..probe_parity.wb import DEFAULT_AS_SHOT_TEMP, DEFAULT_AS_SHOT_TINT, build_wb_matrix
from .color_torch import linear_to_srgb, prophoto_to_srgb, srgb_to_linear, srgb_to_prophoto
from .device import get_torch_device, hwc_numpy_to_torch, hwc_torch_to_numpy
from .lut_torch import apply_polarr_rgb_lut_prophoto_torch

PROPHOTO_LUMA = (0.242655, 0.755158, 0.002187)
T_A, T_B, T_C, T_D, T_E = 1.2, 0.0, 0.96, 0.22, 0.02


def _luma(rgb: torch.Tensor) -> torch.Tensor:
    w = rgb.new_tensor(PROPHOTO_LUMA)
    return (rgb * w).sum(dim=-1)


def _smoothstep(edge0, edge1, x: torch.Tensor) -> torch.Tensor:
    t = torch.clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0)
    return t * t * (3.0 - 2.0 * t)


def _rgb_to_hsv(rgb: torch.Tensor) -> torch.Tensor:
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
    mx = torch.maximum(torch.maximum(r, g), b)
    mn = torch.minimum(torch.minimum(r, g), b)
    d = mx - mn
    h = torch.zeros_like(mx)
    safe_d = torch.where(d > 1e-8, d, torch.ones_like(d))
    h = torch.where((mx == r) & (d > 1e-8), ((g - b) / safe_d + torch.where(g < b, 6.0, 0.0)) / 6.0, h)
    h = torch.where((mx == g) & (d > 1e-8), ((b - r) / safe_d + 2.0) / 6.0, h)
    h = torch.where((mx == b) & (d > 1e-8), ((r - g) / safe_d + 4.0) / 6.0, h)
    h = h % 1.0
    s = torch.where(mx <= 0, torch.zeros_like(mx), d / torch.maximum(mx, torch.full_like(mx, 1e-8)))
    return torch.stack([h, s, mx], dim=-1)


def _hsv_to_rgb(hsv: torch.Tensor) -> torch.Tensor:
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    i = (torch.floor(h * 6.0).to(torch.int64) % 6).to(torch.int64)
    f = h * 6.0 - torch.floor(h * 6.0)
    p = v * (1.0 - s)
    q = v * (1.0 - f * s)
    t = v * (1.0 - (1.0 - f) * s)
    out = torch.zeros_like(hsv)
    for idx, (c0, c1, c2) in enumerate([(v, t, p), (q, v, p), (p, v, t), (p, q, v), (t, p, v), (v, p, q)]):
        m = i == idx
        out[..., 0] = torch.where(m, c0, out[..., 0])
        out[..., 1] = torch.where(m, c1, out[..., 1])
        out[..., 2] = torch.where(m, c2, out[..., 2])
    return out


def _set_hue(rgb: torch.Tensor, hue: torch.Tensor) -> torch.Tensor:
    hsv = _rgb_to_hsv(rgb)
    hsv[..., 0] = hue
    return _hsv_to_rgb(hsv)


def _rgb_to_hue(rgb: torch.Tensor) -> torch.Tensor:
    return _rgb_to_hsv(rgb)[..., 0]


def _tonemap_inv(rgb: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(rgb)
    for ch in range(3):
        x = rgb[..., ch]
        a = x * T_C - T_A
        b_coef = x * T_D - T_B
        c = x * T_E
        disc = torch.clamp(b_coef * b_coef - 4.0 * a * c, min=0.0)
        out[..., ch] = (-b_coef - torch.sqrt(disc)) / (2.0 * a + 1e-12)
    return out


def _tonemap(rgb: torch.Tensor) -> torch.Tensor:
    x = torch.clamp(rgb, 0.0, 32.0)
    return (x * (T_A * x + T_B)) / (x * (T_C * x + T_D) + T_E)


def _apply_wb(rgb: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
    return torch.clamp(rgb @ matrix.T, min=0.0)


def _apply_shadows_highlights(
    rgb: torch.Tensor, shadows: float, highlights: float, luma_map: torch.Tensor
) -> torch.Tensor:
    color_luma = _luma(rgb)
    mixed = torch.clamp(color_luma + luma_map, min=0.0) * 0.5
    luma = _smoothstep(0.0, 1.6, torch.sqrt(mixed))
    specular = torch.clamp(1.0 - luma * luma, min=0.0)
    hi_mask = _smoothstep(0.05, 1.0, luma_map) * 1.2
    if highlights > 0.0:
        hi_mask = hi_mask * torch.clamp(1.0 - color_luma, min=0.0) * 1.1
    sh_mask = 1.0 - _smoothstep(0.0, 0.5, luma_map)
    exposure_amount = shadows * sh_mask + highlights * hi_mask * specular
    return rgb * (2.0 ** (exposure_amount * 2.0)).unsqueeze(-1)


def _apply_saturation(rgb: torch.Tensor, saturation: float) -> torch.Tensor:
    luma = _luma(rgb).unsqueeze(-1)
    mix_amount = min(-saturation, 1.0)
    inv = 1.0 - mix_amount
    return torch.clamp(rgb * inv + luma * mix_amount, min=0.0)


def _contract_blacks_whites(rgb: torch.Tensor, whites: float, blacks: float) -> torch.Tensor:
    blacks = max(blacks, 0.0)
    whites = min(whites, 0.0)
    luma = _luma(rgb)
    x = torch.clamp(torch.sqrt(luma), 0.0, 1.0)
    whites_mask = _smoothstep(0.05, 1.0, x)
    specular = torch.clamp(1.0 - luma * luma * luma * luma, min=0.0)
    whites_mask = whites_mask * specular
    blacks_mask = _smoothstep(0.8, -0.4, x)
    mult = (2.0 ** (whites * whites_mask + blacks * blacks_mask)).unsqueeze(-1)
    return rgb * mult


def _falloff(x: torch.Tensor, m: float) -> torch.Tensor:
    xm = x * m
    xm = 1.0 / ((xm + 1.0) ** 2)
    mm = 1.0 / ((m + 1.0) ** 2)
    return (xm - mm) / (1.0 - mm)


def _expand_whites_blacks(rgb: torch.Tensor, whites: float, blacks: float) -> torch.Tensor:
    blacks = min(blacks, 0.0)
    whites = max(whites, 0.0)
    x = torch.clamp(_luma(rgb), 0.0, 1.0)
    whites_mask = _falloff(x, -0.7)
    whites_mask_low = _smoothstep(0.0, 1.0, torch.sqrt(x))
    blacks_mask = _falloff(1.0 - x, -0.9)
    mask = whites_mask * blacks_mask
    blacks_v = blacks * 0.33 * mask
    whites_v = 1.0 + whites * 0.33 * mask * whites_mask_low
    out = (rgb * (1.0 - blacks_v.unsqueeze(-1)) + blacks_v.unsqueeze(-1)) * whites_v.unsqueeze(-1)
    return torch.clamp(out, 0.0, 1.0)


def _contrast_curve(contrast: float, width: int, device: torch.device) -> torch.Tensor:
    if contrast < 0:
        pts = [(0, 0), (8, 21), (16, 34), (32, 54), (64, 88), (96, 115), (128, 135), (160, 152), (192, 174), (224, 206), (240, 228), (255, 255)]
    else:
        pts = [(0, 0), (32, 8), (64, 35), (96, 71), (128, 117), (160, 171), (192, 212), (224, 241), (255, 255)]
    intensity = abs(contrast)
    xs = torch.tensor([p[0] / 255.0 for p in pts], dtype=torch.float64, device=device)
    ys = torch.tensor(
        [p[1] / 255.0 + intensity * (p[1] / 255.0 - p[0] / 255.0) * (1.0 - p[0] / 255.0) for p in pts],
        dtype=torch.float64,
        device=device,
    )
    xs = torch.clamp(xs, 0, 1)
    ys = torch.clamp(ys, 0, 1)
    samples = torch.linspace(0, 1, width, device=device, dtype=torch.float32)
    return torch.from_numpy(
        __import__("numpy").interp(samples.cpu().numpy(), xs.cpu().numpy(), ys.cpu().numpy()).astype("float32")
    ).to(device)


def _apply_curve(rgb: torch.Tensor, curve: torch.Tensor) -> torch.Tensor:
    rgb = torch.clamp(rgb, 0.0, 1.0)
    size = curve.shape[0]
    out = torch.empty_like(rgb)
    for ch in range(3):
        values = rgb[..., ch]
        idx = torch.clamp(torch.floor(values * (size - 1)).to(torch.int64), 0, size - 2)
        p0x = idx.float() / (size - 1)
        p1x = (idx + 1).float() / (size - 1)
        t = (values - p0x) / torch.clamp(p1x - p0x, min=1e-12)
        out[..., ch] = torch.clamp(curve[idx] + t * (curve[idx + 1] - curve[idx]), 0.0, 1.0)
    return out


def _sample_bilinear_table(
    data: torch.Tensor, width: int, height: int, u: torch.Tensor, v: torch.Tensor
) -> torch.Tensor:
    x = torch.clamp(u, 0.0, 1.0) * (width - 1)
    y = torch.clamp(v, 0.0, 1.0) * (height - 1)
    x0 = torch.floor(x).to(torch.int64)
    y0 = torch.floor(y).to(torch.int64)
    x1 = torch.minimum(x0 + 1, torch.tensor(width - 1, device=data.device))
    y1 = torch.minimum(y0 + 1, torch.tensor(height - 1, device=data.device))
    fx = (x - x0.float()).unsqueeze(-1)
    fy = (y - y0.float()).unsqueeze(-1)
    flat = data.reshape(height, width, 3)
    c00, c01 = flat[y0, x0], flat[y0, x1]
    c10, c11 = flat[y1, x0], flat[y1, x1]
    c0 = c00 * (1.0 - fx) + c01 * fx
    c1 = c10 * (1.0 - fx) + c11 * fx
    return c0 * (1.0 - fy) + c1 * fy


def _apply_hsv_look_table(
    rgb: torch.Tensor, dims: tuple[int, int, int], table: torch.Tensor
) -> torch.Tensor:
    hue_divs, sat_divs, val_divs = dims
    inv_sat = 1.0 / sat_divs
    inv_val = 1.0 / val_divs
    hsv = _rgb_to_hsv(rgb)
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    x = (s * (sat_divs - 1.0) + 0.5) * inv_sat
    y = (h * hue_divs + 0.5) / hue_divs
    z = v * (val_divs - 1.0)
    z_floor = torch.floor(z)
    zf = (z_floor + y) * inv_val
    zc = (torch.ceil(z) + y) * inv_val
    width = sat_divs
    height = hue_divs * val_divs
    tex = table.reshape(height, width, 3)
    col1 = _sample_bilinear_table(tex, width, height, x, zf)
    col2 = _sample_bilinear_table(tex, width, height, x, zc)
    fract_z = (z - z_floor).unsqueeze(-1)
    factors = col1 * (1.0 - fract_z) + col2 * fract_z
    hsv[..., 0] = (h + factors[..., 0]) % 1.0
    hsv[..., 1] = s * factors[..., 1]
    hsv[..., 2] = v * factors[..., 2]
    return _hsv_to_rgb(hsv)


@torch.inference_mode()
def render_srgb_develop_torch(
    srgb_hwc: torch.Tensor,
    adjustments: dict,
    *,
    merged_lut=None,
    lut_strength: float = 1.0,
    lut_encoding: str = "srgb_srgb",
    force_color_look: bool = False,
    input_tone_map_inversed: bool = False,
    as_shot_temp: float = DEFAULT_AS_SHOT_TEMP,
    as_shot_tint: float = DEFAULT_AS_SHOT_TINT,
) -> torch.Tensor:
    """Develop + optional LUT on GPU; input/output H×W×3 float [0,1] on device."""
    shape = srgb_hwc.shape
    device = srgb_hwc.device
    if input_tone_map_inversed:
        rgb = torch.clamp(srgb_hwc, 0.0, 1.0).reshape(-1, 3)
    else:
        rgb = srgb_to_prophoto(srgb_hwc).reshape(-1, 3)
        hue0 = _rgb_to_hue(rgb)
        rgb = _tonemap_inv(rgb)
        rgb = _set_hue(rgb, hue0)
        rgb = torch.clamp(rgb, min=0.0)

    gamma = linear_to_srgb(rgb)
    luma_map = _luma(gamma)

    wb_np = build_wb_matrix(
        adjustments.get("temperature", as_shot_temp),
        adjustments.get("tint", as_shot_tint),
        as_shot_temp=as_shot_temp,
        as_shot_tint=as_shot_tint,
    )
    wb = torch.tensor(wb_np, dtype=torch.float32, device=device)
    rgb = _apply_wb(rgb, wb)

    exposure_factor = float(2.0 ** adjustments.get("exposure", 0.0))
    shadows = float(adjustments.get("shadows", 0.0))
    highlights = float(adjustments.get("highlights", 0.0))
    whites = float(adjustments.get("whites", 0.0))
    blacks = float(adjustments.get("blacks", 0.0))
    saturation = float(adjustments.get("saturation", 0.0))
    contrast = float(adjustments.get("contrast", 0.0))

    if exposure_factor != 1.0:
        rgb = rgb * exposure_factor
    if shadows != 0.0 or highlights != 0.0:
        rgb = _apply_shadows_highlights(rgb, shadows, highlights, luma_map)
    if whites < 0.0 or blacks > 0.0:
        rgb = _contract_blacks_whites(rgb, whites, blacks)

    hue = _rgb_to_hue(rgb)
    rgb = torch.clamp(rgb, 0.0, 1.0)
    rgb = _set_hue(rgb, hue)
    if saturation != 0.0:
        rgb = _apply_saturation(rgb, saturation)

    if force_color_look:
        from ..probe_parity.adobe_assets import load_adobe_profile_look_table, profile_look_dims

        dims = profile_look_dims()
        table = torch.from_numpy(load_adobe_profile_look_table()).to(device=device)
        rgb = _apply_hsv_look_table(rgb, dims, table)

    hue = _rgb_to_hue(rgb)
    rgb = _tonemap(rgb)
    rgb = _set_hue(rgb, hue)
    rgb = linear_to_srgb(rgb)
    hue = _rgb_to_hue(rgb)
    if whites > 0.0 or blacks < 0.0:
        rgb = _expand_whites_blacks(rgb, whites, blacks)
    if contrast != 0.0:
        rgb = _apply_curve(rgb, _contrast_curve(contrast, 256, device))
    rgb = _set_hue(rgb, hue)
    rgb = srgb_to_linear(rgb)

    if merged_lut is not None and lut_strength > 0.0:
        gamma_id, primaries_id = ENCODING_PRESETS.get(lut_encoding, ENCODING_PRESETS["srgb_srgb"])
        rgb = apply_polarr_rgb_lut_prophoto_torch(
            rgb.reshape(shape),
            merged_lut,
            lut_strength,
            rgb_gamma=gamma_id,
            rgb_primaries=primaries_id,
        ).reshape(-1, 3)

    return prophoto_to_srgb(rgb.reshape(shape))


def render_srgb_develop_numpy(
    srgb_hwc,
    adjustments: dict,
    *,
    device: torch.device | None = None,
    **kwargs,
):
    """NumPy in/out wrapper for pipeline integration."""
    device = device or get_torch_device()
    t = hwc_numpy_to_torch(srgb_hwc, device)
    out = render_srgb_develop_torch(t, adjustments, **kwargs)
    return hwc_torch_to_numpy(out)