"""RAW decode + camera-space white balance from an illuminant chromaticity.

White balance is applied the way the iOS app does it (`CIRAWFilter.neutralChromaticity`):
as per-channel multipliers in the camera's *native* space, derived from the target
illuminant through the camera's own XYZ->camera matrix. The earlier Bradford-CAT-in-sRGB
approach used the wrong operator in the wrong space and produced a desaturated,
magenta/hazy result; this neutralises the illuminant the way the app does.

This is deliberately a *white-balance* preview, not a full ISP: no auto-exposure, tone
curve, or contrast. 'before' and 'after' share identical decode settings and differ
only in the WB multipliers, so the comparison shows the correction alone.
"""

import numpy as np

HALF_SIZE = True
MAX_LONG_EDGE = 2000

# Shared postprocess: linear 16-bit, no auto-exposure, half-res. WB + output colour
# are supplied per call so before/after differ only by white balance.
_POST = dict(output_bps=16, gamma=(1, 1), no_auto_bright=True, half_size=HALF_SIZE)


def xy_to_XYZ(x, y, Y=1.0):
    return np.array([Y * x / y, Y, Y * (1.0 - x - y) / y])


def open_raw(path):
    """imread and return the RawPy object. The caller may delete the file right
    after — the unpacked sensor data stays in RAM, so depth re-runs re-render from
    memory (no second read) and the upload is still discarded for privacy."""
    import rawpy

    return rawpy.imread(str(path))


def wb_from_chromaticity(raw, illuminant_xy):
    """Raw WB multipliers [R, G, B, G2] that render `illuminant_xy` neutral, derived
    from the camera's XYZ->camera matrix (`raw.rgb_xyz_matrix`). This is the camera-
    space equivalent of CIRAWFilter.neutralChromaticity. Validated: feeding D65 here
    reproduces `raw.daylight_whitebalance` to 4 decimals."""
    M = np.asarray(raw.rgb_xyz_matrix, dtype=float)[:3]   # XYZ -> camera
    cam = M @ xy_to_XYZ(*illuminant_xy)
    cam = np.where(np.abs(cam) < 1e-9, 1e-9, cam)
    mul = 1.0 / cam
    mul = mul / mul[1]                                    # normalise green = 1
    return [float(mul[0]), 1.0, float(mul[2]), 1.0]


def baseline_wb(raw):
    """'before' WB — the camera's as-shot balance (the photo as captured), falling
    back to the daylight balance. Green-normalised."""
    cw = np.asarray(raw.camera_whitebalance, dtype=float)
    if cw.size < 3 or cw[1] <= 0 or not np.all(np.isfinite(cw[:3])):
        cw = np.asarray(raw.daylight_whitebalance, dtype=float)
    cw = cw / cw[1]
    return [float(cw[0]), 1.0, float(cw[2]), 1.0]


def render(raw, wb, max_long_edge=MAX_LONG_EDGE):
    """postprocess `raw` with WB multipliers `wb` -> uint8 sRGB display image."""
    import rawpy

    rgb16 = raw.postprocess(user_wb=list(wb), output_color=rawpy.ColorSpace.sRGB, **_POST)
    linear = _resize_max(rgb16.astype(np.float32) / 65535.0, max_long_edge)
    return _encode_display(linear)


def _encode_display(linear_rgb):
    lin = np.clip(linear_rgb, 0.0, 1.0)
    srgb = np.where(lin <= 0.0031308, lin * 12.92, 1.055 * np.power(lin, 1 / 2.4) - 0.055)
    return (np.clip(srgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)


def _resize_max(img, max_long_edge):
    from PIL import Image

    h, w = img.shape[:2]
    if max(h, w) <= max_long_edge:
        return img
    scale = max_long_edge / max(h, w)
    new_w, new_h = max(1, round(w * scale)), max(1, round(h * scale))
    chans = [
        np.asarray(Image.fromarray(img[..., c], mode="F").resize((new_w, new_h), Image.BILINEAR))
        for c in range(img.shape[2])
    ]
    return np.stack(chans, axis=-1)