#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "numpy",
#     "scipy",
#     "pillow",
#     "openexr",
#     "imath",
#     "matplotlib"
# ]
# ///

"""
Gaussianization for Histogram-Preserving Blending
Based on Burley's "On Histogram-preserving Blending for Randomized Texture Tiling" (2019)

This implementation uses per-channel 1D histogram transformation with:
- Truncated Gaussian distribution (avoiding values outside [0,1])
- Soft-clipping contrast operator
- Fast 1D LUT generation from input histogram
"""

import argparse
import sys
from pathlib import Path
import numpy as np
from scipy import special
from PIL import Image
import OpenEXR
import Imath


def load_image(image_path: Path) -> np.ndarray:
    """Load a PNG or EXR image and return as float64 array (H, W, 3) in [0,1]."""
    suffix = image_path.suffix.lower()
    if suffix == '.exr':
        exr = OpenEXR.InputFile(str(image_path))
        header = exr.header()
        dw = header['dataWindow']
        w = dw.max.x - dw.min.x + 1
        h = dw.max.y - dw.min.y + 1
        channels = []
        for ch in ('R', 'G', 'B'):
            raw = exr.channel(ch, Imath.PixelType(Imath.PixelType.HALF))
            channels.append(np.frombuffer(raw, dtype=np.float16).reshape(h, w))
        return np.stack(channels, axis=-1).astype(np.float64)
    else:
        return np.array(Image.open(image_path).convert("RGB")).astype(np.float64) / 255.0


def save_image(image: np.ndarray, output_path: Path):
    """Save an RGB float image as EXR or PNG depending on the file suffix."""
    suffix = output_path.suffix.lower()
    if suffix == ".exr":
        image_to_save = image.astype(np.float16)
        h, w, _ = image_to_save.shape
        header = OpenEXR.Header(w, h)
        half_chan = Imath.Channel(Imath.PixelType(Imath.PixelType.HALF))
        header['channels'] = {'R': half_chan, 'G': half_chan, 'B': half_chan}
        out = OpenEXR.OutputFile(str(output_path), header)
        out.writePixels({
            'R': image_to_save[:, :, 0].tobytes(),
            'G': image_to_save[:, :, 1].tobytes(),
            'B': image_to_save[:, :, 2].tobytes(),
        })
        out.close()
    elif suffix == ".png":
        clipped = np.clip(image, 0.0, 1.0)
        Image.fromarray(np.round(clipped * 255.0).astype(np.uint8), mode="RGB").save(output_path)
    else:
        raise ValueError(f"Unsupported output format '{output_path.suffix}'. Use .exr or .png.")


class TruncatedGaussian:
    """Truncated Gaussian distribution for histogram transformation."""

    def __init__(self, sigma: float = 1.0 / 6.0):
        """
        Initialize truncated Gaussian centered at 0.5 with given sigma.
        Default sigma = 1/6 as recommended in Burley's paper.
        """
        self.sigma = sigma
        self.mu = 0.5

        # Burley's C(sigma) is the reciprocal normalization factor required
        # after truncating the Gaussian to the [0, 1] interval.
        self.C = 1.0 / special.erf(1.0 / (2.0 * np.sqrt(2.0) * sigma))

    def inverse_cdf(self, u: np.ndarray) -> np.ndarray:
        """
        Inverse CDF of truncated Gaussian distribution.
        Maps uniform values in [0,1] to truncated Gaussian in [0,1].

        Equation (3) from Burley's paper:
        CDF^-1_[G](u; σ) = 1/2 + sqrt(2)σ * erfinv((2u - 1) / C(σ))
        """
        u = np.clip(u, 0.0, 1.0)
        result = 0.5 + np.sqrt(2.0) * self.sigma * special.erfinv((2.0 * u - 1.0) / self.C)
        return np.clip(result, 0.0, 1.0)

    def cdf(self, x: np.ndarray) -> np.ndarray:
        """
        CDF of truncated Gaussian distribution.
        Maps truncated Gaussian values to uniform [0,1].
        """
        x = np.clip(x, 0.0, 1.0)
        result = 0.5 * (1.0 + self.C * special.erf((x - 0.5) / (np.sqrt(2.0) * self.sigma)))
        return np.clip(result, 0.0, 1.0)


def _cdf_bin_edges(histogram: np.ndarray) -> np.ndarray:
    """Return normalized CDF samples on histogram bin edges."""
    histogram = np.asarray(histogram, dtype=np.float64)
    total = float(histogram.sum())
    if total <= 0.0:
        return np.linspace(0.0, 1.0, len(histogram) + 1)
    return np.concatenate(([0.0], np.cumsum(histogram / total, dtype=np.float64)))


def _occupied_bin_mid_quantiles(histogram: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Return source-value centers and CDF midpoints for occupied bins."""
    histogram = np.asarray(histogram, dtype=np.float64)
    n_bins = len(histogram)
    cdf_edges = _cdf_bin_edges(histogram)
    bin_mass = cdf_edges[1:] - cdf_edges[:-1]
    occupied = bin_mass > 0.0

    if not np.any(occupied):
        centers = (np.arange(n_bins, dtype=np.float64) + 0.5) / n_bins
        return centers, centers

    centers = (np.arange(n_bins, dtype=np.float64) + 0.5) / n_bins
    mid_quantiles = cdf_edges[:-1] + 0.5 * bin_mass
    return centers[occupied], mid_quantiles[occupied]


def build_gaussianization_lut(histogram: np.ndarray, lut_size: int = 4096) -> np.ndarray:
    """
    Build 1D LUT for Gaussianizing a channel based on its histogram.

    Algorithm 1 from Burley's paper:
    1. Compute CDF from histogram
    2. Transform through inverse CDF of truncated Gaussian
    """
    gaussian = TruncatedGaussian()
    value_centers, mid_quantiles = _occupied_bin_mid_quantiles(histogram)
    mapped_centers = gaussian.inverse_cdf(mid_quantiles)

    sample_positions = np.linspace(0.0, 1.0, lut_size)
    return np.interp(
        sample_positions,
        value_centers,
        mapped_centers,
        left=mapped_centers[0],
        right=mapped_centers[-1],
    )


def apply_lut(image: np.ndarray, lut: np.ndarray) -> np.ndarray:
    """Apply a 1D LUT to an image channel using linear interpolation."""
    lut_size = len(lut)
    coords = np.clip(image, 0.0, 1.0) * (lut_size - 1)
    indices0 = np.floor(coords).astype(np.int32)
    indices1 = np.minimum(indices0 + 1, lut_size - 1)
    alpha = coords - indices0
    return (1.0 - alpha) * lut[indices0] + alpha * lut[indices1]


def _deterministic_noise(shape: tuple[int, int], channel: int) -> np.ndarray:
    """Generate stable per-pixel noise in [0, 1) from pixel coordinates."""
    height, width = shape
    yy, xx = np.indices((height, width), dtype=np.uint32)
    state = xx * np.uint32(0x1F123BB5) ^ yy * np.uint32(0x159A55E5) ^ np.uint32(channel + 1) * np.uint32(0x2C1B3C6D)
    state ^= state >> np.uint32(16)
    state *= np.uint32(0x7FEB352D)
    state ^= state >> np.uint32(15)
    state *= np.uint32(0x846CA68B)
    state ^= state >> np.uint32(16)
    return state.astype(np.float64) / float(np.iinfo(np.uint32).max)


def dither_channel(channel: np.ndarray, quantization_step: float, channel_index: int) -> np.ndarray:
    """Spread repeated quantized values across their source bucket deterministically."""
    if quantization_step <= 0.0:
        return channel
    noise = _deterministic_noise(channel.shape, channel_index) - 0.5
    return np.clip(channel + noise * quantization_step, 0.0, 1.0)


def _soft_clipping_lower_half(x_hat: np.ndarray, W_hat: float) -> np.ndarray:
    """Evaluate Burley's Eq. 4 on the lower half of the domain."""
    linear_start = (2.0 - W_hat) / 4.0
    linear = (x_hat - 0.5) / W_hat + 0.5

    if W_hat >= (2.0 / 3.0):
        t = x_hat / (2.0 - W_hat)
        quadratic = 8.0 * (1.0 / W_hat - 1.0) * (t ** 2) + (3.0 - 2.0 / W_hat) * t
        return np.where(x_hat >= linear_start, linear, quadratic)

    quadratic_start = (2.0 - 3.0 * W_hat) / 4.0
    quadratic = ((x_hat - quadratic_start) / W_hat) ** 2
    return np.where(
        x_hat >= linear_start,
        linear,
        np.where(x_hat >= quadratic_start, quadratic, 0.0),
    )


def soft_clipping_contrast(x_hat: np.ndarray, W_hat: float) -> np.ndarray:
    """
    Soft-clipping contrast operator S*_[G] from Equation (4) in Burley's paper.

    This is a piecewise function that:
    - Is linear in the middle half of the range
    - Blends smoothly to 0 or 1 using quadratic segments at the ends
    """
    if not (0.0 < W_hat <= 1.0):
        raise ValueError(f"W_hat must be in (0, 1], got {W_hat}")

    x_hat = np.clip(x_hat, 0.0, 1.0)
    lower_input = np.where(x_hat <= 0.5, x_hat, 1.0 - x_hat)
    lower_result = _soft_clipping_lower_half(lower_input, W_hat)
    result = np.where(x_hat <= 0.5, lower_result, 1.0 - lower_result)
    return np.clip(result, 0.0, 1.0)


def gaussianize_texture(
    image: np.ndarray,
    verbose: bool = True,
    quantization_step: float = 0.0,
) -> tuple[np.ndarray, list]:
    """
    Gaussianize a texture using per-channel 1D histogram transformation.

    Returns:
        - Gaussianized image
        - List of inverse LUTs (one per channel) for restoration
    """
    _, _, c = image.shape
    if c != 3:
        raise ValueError(f"Expected RGB image with 3 channels, got {c}")

    # Process each channel independently
    gaussianized = np.zeros_like(image)
    inverse_luts = []

    for ch in range(3):
        if verbose:
            channel_name = ['R', 'G', 'B'][ch]
            print(f"Processing channel {channel_name}...")

        # Break ties inside quantized source buckets before building the transport.
        channel = dither_channel(image[:, :, ch], quantization_step, ch)

        # Compute histogram (using 4096 bins for better precision)
        hist, _ = np.histogram(channel.flatten(), bins=4096, range=(0.0, 1.0))

        # Build Gaussianization LUT
        lut = build_gaussianization_lut(hist, lut_size=4096)

        # Apply LUT to channel
        gaussianized[:, :, ch] = apply_lut(channel, lut)

        # Build inverse LUT for later restoration
        inverse_lut = build_inverse_lut(hist, lut_size=4096)
        inverse_luts.append(inverse_lut)

    return gaussianized, inverse_luts


def build_inverse_lut(original_histogram: np.ndarray, lut_size: int = 4096) -> np.ndarray:
    """
    Build inverse LUT to restore original histogram from Gaussianized values.
    This maps from Gaussian distribution back to original distribution.
    """
    gaussian = TruncatedGaussian()
    value_centers, mid_quantiles = _occupied_bin_mid_quantiles(original_histogram)

    gaussian_values = np.linspace(0.0, 1.0, lut_size)
    uniform_values = gaussian.cdf(gaussian_values)
    return np.interp(
        uniform_values,
        mid_quantiles,
        value_centers,
        left=value_centers[0],
        right=value_centers[-1],
    )


def histogram_preserving_blend(
    textures: list[np.ndarray],
    weights: np.ndarray,
    inverse_luts: list[np.ndarray] | list[list[np.ndarray]],
    gamma: float = 1.0
) -> np.ndarray:
    """
    Perform histogram-preserving blend of multiple Gaussianized textures.

    Algorithm 2 from Burley's paper:
    1. Optionally exponentiate weights
    2. Linear blend
    3. Compute variance scale factor W_hat
    4. Apply soft-clipping contrast operator
    5. Apply inverse LUTs to restore original histogram

    Args:
        textures: List of Gaussianized textures
        weights: Blending weights (must sum to 1)
        inverse_luts: Shared inverse LUTs for the source texture, or repeated copies
            of the same LUT set for each texture.
        gamma: Exponent for weight adjustment (Eq. 5)
    """
    n_textures = len(textures)
    if len(weights) != n_textures:
        raise ValueError(f"Number of weights ({len(weights)}) must match number of textures ({n_textures})")

    if len(inverse_luts) == 3 and all(np.asarray(lut).ndim == 1 for lut in inverse_luts):
        shared_inverse_luts = inverse_luts
    else:
        if len(inverse_luts) != n_textures:
            raise ValueError(
                "inverse_luts must be either one shared RGB LUT set or one repeated set per texture"
            )
        shared_inverse_luts = inverse_luts[0]
        for lut_set in inverse_luts[1:]:
            if any(not np.array_equal(ref, cur) for ref, cur in zip(shared_inverse_luts, lut_set)):
                raise ValueError(
                    "Burley's per-channel method assumes all blended tiles share the same histogram LUTs"
                )

    # Normalize weights
    weights = np.array(weights, dtype=np.float64) / np.sum(weights)

    # Apply weight exponentiation if gamma != 1 (Equation 5)
    if gamma != 1.0:
        weights_exp = np.power(weights, gamma)
        weights = weights_exp / np.sum(weights_exp)

    # Linear blend (Equation 1)
    blended = np.zeros_like(textures[0])
    for tex, w in zip(textures, weights):
        blended += w * tex

    # Compute variance scale factor (Equation 2)
    W_hat = np.sqrt(np.sum(weights ** 2))

    # Apply contrast restoration per channel
    result = np.zeros_like(blended)
    for ch in range(3):
        # Apply soft-clipping contrast operator (Equation 4)
        result[:, :, ch] = soft_clipping_contrast(blended[:, :, ch], W_hat)

        # Apply inverse LUT to restore the shared source histogram.
        result[:, :, ch] = apply_lut(result[:, :, ch], shared_inverse_luts[ch])

    return result


def verify_histogram(image_path: Path, output_path: Path):
    """Generate a minimal RGB histogram verification figure."""
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    img = load_image(image_path)

    fig = plt.figure(figsize=(4.0, 2.0), facecolor='#bcbcbc')
    plot_ax = fig.add_axes((0.0, 0.0, 1.0, 1.0))
    plot_ax.set_facecolor('#bcbcbc')
    for spine in plot_ax.spines.values():
        spine.set_visible(False)
    plot_ax.set_xticks([])
    plot_ax.set_yticks([])

    kernel = np.array([1.0, 2.0, 3.0, 2.0, 1.0], dtype=np.float64)
    kernel /= kernel.sum()
    curve_max = 0.0
    colors = ('#ff1a1a', '#00aa22', '#003cff')

    for channel, color in enumerate(colors):
        hist, edges = np.histogram(img[:, :, channel].ravel(), bins=512, range=(0.0, 1.0), density=True)
        hist = np.convolve(hist, kernel, mode='same')
        centers = 0.5 * (edges[:-1] + edges[1:])
        curve_max = max(curve_max, float(hist.max()))
        plot_ax.plot(centers, hist, color=color, linewidth=0.9, antialiased=True)

    plot_ax.set_xlim(0.0, 1.0)
    plot_ax.set_ylim(0.0, curve_max * 1.18 if curve_max > 0.0 else 1.0)
    fig.savefig(output_path, dpi=180, facecolor=fig.get_facecolor(), edgecolor='none')
    plt.close(fig)
    print(f"Saved histogram to {output_path}")


def save_lut_as_image(
    luts: list[np.ndarray],
    output_path: Path,
    width: int = 2048,
    height: int = 2048,
):
    """Save inverse LUTs as a 2D texture with LUT samples running across columns."""
    if width <= 0 or height <= 0:
        raise ValueError(f"Invalid LUT image size {width}x{height}")

    lut_size = len(luts[0])
    if any(len(lut) != lut_size for lut in luts):
        raise ValueError("All inverse LUT channels must have the same length")

    src_coords = np.linspace(0.0, 1.0, lut_size)
    dst_coords = np.linspace(0.0, 1.0, width)
    packed_columns = np.stack(
        [np.interp(dst_coords, src_coords, lut) for lut in luts],
        axis=-1,
    )

    lut_image = np.broadcast_to(packed_columns[np.newaxis, :, :], (height, width, 3)).copy()
    save_image(lut_image, output_path)


def main():
    parser = argparse.ArgumentParser(
        description="Gaussianize texture using Burley's per-channel histogram-preserving method",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        "input",
        type=Path,
        help="Path to the input texture (PNG or EXR)"
    )
    parser.add_argument(
        "-o", "--output",
        type=Path,
        default=None,
        help="Output path. Defaults to <input>_gaussianized.exr"
    )
    parser.add_argument(
        "--inverse-lut",
        action="store_true",
        default=True,
        help="Also save the inverse LUT as <input>_inverse_lut.exr"
    )
    parser.add_argument(
        "--verify",
        action="store_true",
        help="Generate histogram visualization instead of processing"
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Print detailed progress information"
    )

    args = parser.parse_args()

    if not args.input.exists():
        print(f"Error: Input file '{args.input}' does not exist.", file=sys.stderr)
        sys.exit(1)

    if args.verify:
        # Generate histogram visualization
        hist_path = args.input.with_name(args.input.stem + "_histogram.png")
        verify_histogram(args.input, hist_path)
    else:
        # Load input image
        print(f"Loading {args.input}...")
        image = load_image(args.input)
        quantization_step = 0.0 if args.input.suffix.lower() == ".exr" else (1.0 / 255.0)

        # Gaussianize the texture
        print("Applying per-channel Gaussianization...")
        gaussianized, inverse_luts = gaussianize_texture(
            image,
            verbose=args.verbose,
            quantization_step=quantization_step,
        )

        # Determine output path
        if args.output is None:
            args.output = args.input.with_name(args.input.stem + "_gaussianized.exr")

        # Save Gaussianized texture
        print(f"Saving Gaussianized texture to {args.output}...")
        save_image(gaussianized, args.output)

        # Optionally save inverse LUT
        if args.inverse_lut:
            lut_path = args.input.with_name(args.input.stem + "_inverse_lut.exr")
            print(f"Saving inverse LUT to {lut_path}...")
            save_lut_as_image(inverse_luts, lut_path)

        print("Done!")


if __name__ == "__main__":
    main()