#!/usr/bin/env -S uv run --script
# /// script
# dependencies = [
#   "numpy",
#   "scipy",
#   "Pillow",
#   "imageio",
#   "rawpy",
# ]
# ///

# -*- coding: utf-8 -*-
"""
Single-file Python script for Film Stock Color Emulation based on Datasheet CSV.

Focuses on color transformation, applying effects derived from datasheet parameters.
Assumes input image is in linear RGB format. Excludes film grain simulation.

Dependencies: numpy, imageio, scipy
Installation: pip install numpy imageio scipy Pillow
"""

import argparse
import csv
import numpy as np
import imageio.v3 as iio
from scipy.interpolate import interp1d
from scipy.ndimage import gaussian_filter
import rawpy
import os
import sys
import json

# --- Configuration ---
# Small epsilon to prevent log(0) or division by zero errors
EPSILON = 1e-10

from typing import Optional, List


class Info:
    name: str
    description: str
    format_mm: int
    version: str

    def __init__(
        self, name: str, description: str, format_mm: int, version: str
    ) -> None:
        self.name = name
        self.description = description
        self.format_mm = format_mm
        self.version = version

    def __repr__(self) -> str:
        return (
            f"Info(name={self.name}, description={self.description}, "
            f"format_mm={self.format_mm}, version={self.version})"
        )


class Balance:
    r_shift: float
    g_shift: float
    b_shift: float

    def __init__(self, r_shift: float, g_shift: float, b_shift: float) -> None:
        self.r_shift = r_shift
        self.g_shift = g_shift
        self.b_shift = b_shift
    
    def __repr__(self) -> str:
        return (
            f"Balance(r_shift={self.r_shift:.3f}, g_shift={self.g_shift:.3f}, b_shift={self.b_shift:.3f})"
        )


class Gamma:
    r_factor: float
    g_factor: float
    b_factor: float

    def __init__(self, r_factor: float, g_factor: float, b_factor: float) -> None:
        self.r_factor = r_factor
        self.g_factor = g_factor
        self.b_factor = b_factor

    def __repr__(self) -> str:
        return (
            f"Gamma(r_factor={self.r_factor:.3f}, g_factor={self.g_factor:.3f}, b_factor={self.b_factor:.3f})"
        )


class Processing:
    gamma: Gamma
    balance: Balance

    def __init__(self, gamma: Gamma, balance: Balance) -> None:
        self.gamma = gamma
        self.balance = balance

    def __repr__(self) -> str:
        return (
            f"Processing(gamma=({self.gamma.r_factor:.3f}, {self.gamma.g_factor:.3f}, {self.gamma.b_factor:.3f}), "
            f"balance=({self.balance.r_shift:.3f}, {self.balance.g_shift:.3f}, {self.balance.b_shift:.3f}))"
        )


class Couplers:
    amount: float
    diffusion_um: float

    def __init__(self, amount: float, diffusion_um: float) -> None:
        self.amount = amount
        self.diffusion_um = diffusion_um

    def __repr__(self) -> str:
        return f"Couplers(amount={self.amount:.3f}, diffusion_um={self.diffusion_um:.1f})"


class HDCurvePoint:
    d: Optional[float]
    r: float
    g: float
    b: float

    def __init__(self, d: Optional[float], r: float, g: float, b: float) -> None:
        self.d = d
        self.r = r
        self.g = g
        self.b = b

    def __repr__(self) -> str:
        return f"HDCurvePoint(d={self.d}, r={self.r:.3f}, g={self.g:.3f}, b={self.b:.3f})"
    

class SpectralSensitivityCurvePoint:
    wavelength: float
    y: float
    m: float
    c: float

    def __init__(self, wavelength: float, y: float, m: float, c: float) -> None:
        self.wavelength = wavelength
        self.y = y
        self.m = m
        self.c = c

    def __repr__(self) -> str:
        return f"SpectralSensitivityCurvePoint(wavelength={self.wavelength:.1f}, y={self.y:.3f}, m={self.m:.3f}, c={self.c:.3f})"


class RGBValue:
    r: float
    g: float
    b: float

    def __init__(self, r: float, g: float, b: float) -> None:
        self.r = r
        self.g = g
        self.b = b

    def __repr__(self) -> str:
        return f"RGBValue(r={self.r:.3f}, g={self.g:.3f}, b={self.b:.3f})"


class Curves:
    hd: List[HDCurvePoint]
    spectral_sensitivity: List[SpectralSensitivityCurvePoint]

    def __init__(self, hd: List[HDCurvePoint], spectral_sensitivity: List[SpectralSensitivityCurvePoint]) -> None:
        self.hd = hd
        self.spectral_sensitivity = spectral_sensitivity

    def __repr__(self) -> str:
        return f"Curves(hd={',\n'.join(repr(point) for point in self.hd)}, spectral_sensitivity={',\n'.join(repr(point) for point in self.spectral_sensitivity)})"


class Halation:
    strength: RGBValue
    size_um: RGBValue

    def __init__(self, strength: RGBValue, size_um: RGBValue) -> None:
        self.strength = strength
        self.size_um = size_um

    def __repr__(self) -> str:
        return f"Halation(strength={self.strength}, size_um={self.size_um})"


class Interlayer:
    diffusion_um: float

    def __init__(self, diffusion_um: float) -> None:
        self.diffusion_um = diffusion_um

    def __repr__(self) -> str:
        return f"Interlayer(diffusion_um={self.diffusion_um:.1f})"


class Calibration:
    iso: int
    middle_gray_logE: float

    def __init__(self, iso: int, middle_gray_logE: float) -> None:
        self.iso = iso
        self.middle_gray_logE = middle_gray_logE

    def __repr__(self) -> str:
        return (
            f"Calibration(iso={self.iso}\nmiddle_gray_logE={self.middle_gray_logE:.3f})"
        )


class Properties:
    halation: Halation
    couplers: Couplers
    interlayer: Interlayer
    curves: Curves
    calibration: Calibration

    def __init__(
        self,
        halation: Halation,
        couplers: Couplers,
        interlayer: Interlayer,
        curves: Curves,
        calibration: Calibration,
    ) -> None:
        self.halation = halation
        self.couplers = couplers
        self.interlayer = interlayer
        self.curves = curves
        self.calibration = calibration

    def __repr__(self) -> str:
        return (
            f"Properties(halation={self.halation}\ncouplers={self.couplers}\n"
            f"interlayer={self.interlayer}\ncurves={self.curves}\n"
            f"calibration={self.calibration})"
        )


class FilmDatasheet:
    info: Info
    processing: Processing
    properties: Properties

    def __init__(
        self, info: Info, processing: Processing, properties: Properties
    ) -> None:
        self.info = info
        self.processing = processing
        self.properties = properties

    def __repr__(self) -> str:
        return (
            f"FilmDatasheet(info={self.info}\nprocessing={self.processing}\n"
            f"properties={self.properties})"
        )

import pprint

def parse_datasheet_json(json_filepath) -> FilmDatasheet | None:
    # Parse JSON into FilmDatasheet object
    """Parses the film datasheet JSON file.
    Args:
        json_filepath (str): Path to the datasheet JSON file.
    Returns:
        FilmDatasheet: Parsed datasheet object.
    """
    if not os.path.exists(json_filepath):
        print(f"Error: Datasheet file not found at {json_filepath}", file=sys.stderr)
        return None
    try:
        with open(json_filepath, "r") as jsonfile:
            data = json.load(jsonfile)
            # Parse the JSON data into the FilmDatasheet structure
            info = Info(
                name=data["info"]["name"],
                description=data["info"]["description"],
                format_mm=data["info"]["format_mm"],
                version=data["info"]["version"],
            )
            gamma = Gamma(
                r_factor=data["processing"]["gamma"]["r_factor"],
                g_factor=data["processing"]["gamma"]["g_factor"],
                b_factor=data["processing"]["gamma"]["b_factor"],
            )
            balance = Balance(
                r_shift=data["processing"]["balance"]["r_shift"],
                g_shift=data["processing"]["balance"]["g_shift"],
                b_shift=data["processing"]["balance"]["b_shift"],
            )
            processing = Processing(gamma=gamma, balance=balance)
            halation = Halation(
                strength=RGBValue(
                    r=data["properties"]["halation"]["strength"]["r"],
                    g=data["properties"]["halation"]["strength"]["g"],
                    b=data["properties"]["halation"]["strength"]["b"],
                ),
                size_um=RGBValue(
                    r=data["properties"]["halation"]["size_um"]["r"],
                    g=data["properties"]["halation"]["size_um"]["g"],
                    b=data["properties"]["halation"]["size_um"]["b"],
                ),
            )
            couplers = Couplers(
                amount=data["properties"]["couplers"]["amount"],
                diffusion_um=data["properties"]["couplers"]["diffusion_um"],
            )
            interlayer = Interlayer(
                diffusion_um=data["properties"]["interlayer"]["diffusion_um"]
            )
            calibration = Calibration(
                iso=data["properties"]["calibration"]["iso"],
                middle_gray_logE=data["properties"]["calibration"]["middle_gray_logh"],
            )
            curves = Curves(
                hd=[
                    HDCurvePoint(d=point["d"], r=point["r"], g=point["g"], b=point["b"])
                    for point in data["properties"]["curves"]["hd"]
                ],
                spectral_sensitivity=[
                    SpectralSensitivityCurvePoint(
                        wavelength=point["wavelength"],
                        y=point["y"],
                        m=point["m"],
                        c=point["c"],
                    )
                    for point in data["properties"]["curves"]["spectral_sensitivity"]
                ],
            )
            print(f"Parsed {len(curves.hd)} H&D curve points.")
            properties = Properties(
                calibration=calibration,
                halation=halation,
                couplers=couplers,
                interlayer=interlayer,
                curves=curves,
            )
            return FilmDatasheet(
                info=info, processing=processing, properties=properties
            )
    except Exception as e:
        print(f"Error parsing datasheet JSON '{json_filepath}': {e}", file=sys.stderr)
        return None


def um_to_pixels(sigma_um, image_width_px, film_format_mm):
    """Converts sigma from micrometers to pixels."""
    if film_format_mm <= 0 or image_width_px <= 0:
        return 0
    microns_per_pixel = (film_format_mm * 1000.0) / image_width_px
    sigma_pixels = sigma_um / microns_per_pixel
    return sigma_pixels


def apply_hd_curves(
    log_exposure_rgb,
    processing: Processing,
    hd_curve: List[HDCurvePoint],
    middle_gray_logE: float,
) -> np.ndarray:
    """Applies H&D curves to log exposure values."""
    density_rgb = np.zeros_like(log_exposure_rgb)
    gamma_factors = [
        processing.gamma.r_factor,
        processing.gamma.g_factor,
        processing.gamma.b_factor,
    ]
    balance_shifts = [
        processing.balance.r_shift,
        processing.balance.g_shift,
        processing.balance.b_shift,
    ]

    min_logE = hd_curve[0].d
    max_logE = hd_curve[-1].d
    min_densities = [hd_curve[0].g, hd_curve[0].g, hd_curve[0].g]
    max_densities = [hd_curve[-1].g, hd_curve[-1].g, hd_curve[-1].g]

    for i, channel in enumerate(["R", "G", "B"]):
        # Apply gamma factor (affects contrast by scaling log exposure input)
        # Handle potential division by zero if gamma factor is 0
        gamma_factor = gamma_factors[i]
        if abs(gamma_factor) < EPSILON:
            print(
                f"Warning: Gamma factor for channel {channel} is near zero. Clamping to {EPSILON}.",
                file=sys.stderr,
            )
            gamma_factor = EPSILON if gamma_factor >= 0 else -EPSILON

        # Adjust log exposure relative to middle gray before applying gamma
        log_exposure_adjusted = (
            middle_gray_logE
            + (log_exposure_rgb[..., i] - middle_gray_logE) / gamma_factor
        )

        # Clamp adjusted exposure to the range defined in the H&D data before interpolation
        log_exposure_clamped = np.clip(log_exposure_adjusted, min_logE, max_logE)

        # Create interpolation function for the current channel
        if channel == "R":
            interp_func = interp1d(
                [d.d for d in hd_curve],
                [d.r for d in hd_curve],
                kind="linear",  # Linear interpolation is common, could be 'cubic'
                bounds_error=False,  # Allows extrapolation, but we clamp manually below
                fill_value=(min_densities[i], max_densities[i]),  # type: ignore
            )
        elif channel == "G":
            interp_func = interp1d(
                [d.d for d in hd_curve],
                [d.g for d in hd_curve],
                kind="linear",  # Linear interpolation is common, could be 'cubic'
                bounds_error=False,  # Allows extrapolation, but we clamp manually below
                fill_value=(min_densities[i], max_densities[i]),  # type: ignore
            )
        else:
            interp_func = interp1d(
                [d.d for d in hd_curve],
                [d.b for d in hd_curve],
                kind="linear",  # Linear interpolation is common, could be 'cubic'
                bounds_error=False,  # Allows extrapolation, but we clamp manually below
                fill_value=(min_densities[i], max_densities[i]),  # type: ignore
            )

        # Apply interpolation
        density = interp_func(log_exposure_clamped)

        # Apply density balance shift (additive density offset)
        density += balance_shifts[i]

        density_rgb[..., i] = np.maximum(density, 0)  # Ensure density is non-negative

    return density_rgb


def apply_saturation_rgb(image_linear, saturation_factor):
    """Adjusts saturation directly in RGB space."""
    if saturation_factor == 1.0:
        return image_linear

    # Luminance weights for sRGB primaries (Rec.709)
    luminance = (
        0.2126 * image_linear[..., 0]
        + 0.7152 * image_linear[..., 1]
        + 0.0722 * image_linear[..., 2]
    )

    # Expand luminance to 3 channels for broadcasting
    luminance_rgb = np.expand_dims(luminance, axis=-1)

    # Apply saturation: Lerp between luminance (grayscale) and original color
    saturated_image = luminance_rgb + saturation_factor * (image_linear - luminance_rgb)

    # Clip results to valid range (important after saturation boost)
    return np.clip(saturated_image, 0.0, 1.0)


def apply_spatial_effects(
    image,
    film_format_mm,
    couplerData: Couplers,
    interlayerData: Interlayer,
    halationData: Halation,
    image_width_px,
):
    """Applies diffusion blur and halation."""
    # Combine diffusion effects (assuming they add quadratically in terms of sigma)
    total_diffusion_um = np.sqrt(
        couplerData.diffusion_um**2 + interlayerData.diffusion_um**2
    )

    if total_diffusion_um > EPSILON:
        sigma_pixels_diffusion = um_to_pixels(
            total_diffusion_um, image_width_px, film_format_mm
        )
        if sigma_pixels_diffusion > EPSILON:
            print(
                f"Applying diffusion blur: sigma={sigma_pixels_diffusion:.2f} pixels ({total_diffusion_um:.1f} um)"
            )
            # Apply blur to the linear image data
            image = gaussian_filter(
                image,
                sigma=[sigma_pixels_diffusion, sigma_pixels_diffusion, 0],
                mode="nearest",
            )  # Blur R, G, B independently
            image = np.clip(image, 0.0, 1.0)  # Keep values in range

    # --- 2. Apply Halation ---
    # This simulates light scattering back through the emulsion
    halation_applied = False
    blurred_image_halation = np.copy(
        image
    )  # Start with potentially diffusion-blurred image

    strengths = [
        halationData.strength.r,
        halationData.strength.g,
        halationData.strength.b,
    ]
    sizes_um = [
        halationData.size_um.r,
        halationData.size_um.g,
        halationData.size_um.b,
    ]

    for i in range(3):
        strength = strengths[i]
        size_um = sizes_um[i]
        if strength > EPSILON and size_um > EPSILON:
            sigma_pixels_halation = um_to_pixels(
                size_um, image_width_px, film_format_mm
            )
            if sigma_pixels_halation > EPSILON:
                halation_applied = True
                print(
                    f"Applying halation blur (Channel {i}): sigma={sigma_pixels_halation:.2f} pixels ({size_um:.1f} um), strength={strength:.3f}"
                )
                # Blur only the current channel for halation effect
                channel_blurred = gaussian_filter(
                    image[..., i], sigma=sigma_pixels_halation, mode="nearest"
                )
                # Add the blurred channel back, weighted by strength
                blurred_image_halation[..., i] = (
                    image[..., i] * (1.0 - strength) + channel_blurred * strength
                )

    if halation_applied:
        # Clip final result after halation
        image = np.clip(blurred_image_halation, 0.0, 1.0)

    return image


def main():
    parser = argparse.ArgumentParser(
        description="Simulate film stock color characteristics using a datasheet JSON."
    )
    parser.add_argument(
        "input_image",
        help="Path to the input RGB image (e.g., PNG, TIFF). Assumed linear RGB.",
    )
    parser.add_argument("datasheet_json", help="Path to the film datasheet JSON file.")
    parser.add_argument("output_image", help="Path to save the output emulated image.")

    args = parser.parse_args()

    # --- Load Datasheet ---
    print(f"Loading datasheet: {args.datasheet_json}")
    datasheet: FilmDatasheet | None = parse_datasheet_json(args.datasheet_json)
    if datasheet is None:
        sys.exit(1)
    print(
        f"Simulating: {datasheet.info.name} ({datasheet.info.format_mm}mm) (v{datasheet.info.version})\n\t{datasheet.info.description}"
    )

    import pprint
    pprint.pp(datasheet)

    # --- Load Input Image ---
    print(f"Loading input image: {args.input_image}")
    try:
        # For DNG files, force reading the raw image data, not the embedded thumbnail
        if (
            args.input_image.lower().endswith(".dng")
            or args.input_image.lower().endswith(".raw")
            or args.input_image.lower().endswith(".arw")
        ):
            print("Detected Camera RAW file, reading raw image data...")
            with rawpy.imread(args.input_image) as raw:
                image_raw = raw.postprocess(
                    demosaic_algorithm=rawpy.DemosaicAlgorithm.AHD,  # type: ignore
                    output_bps=16,  # Use 16-bit output for better precision
                    use_camera_wb=True,  # Use camera white balance
                    no_auto_bright=True,  # Disable auto brightness adjustment
                    output_color=rawpy.ColorSpace.sRGB,  # type: ignore
                )
            # If the image has more than 3 channels, try to select the first 3 (RGB)
            if image_raw.ndim == 3 and image_raw.shape[-1] > 3:
                image_raw = image_raw[..., :3]
        else:
            image_raw = iio.imread(args.input_image)
    except FileNotFoundError:
        print(
            f"Error: Input image file not found at {args.input_image}", file=sys.stderr
        )
        sys.exit(1)
    except Exception as e:
        print(f"Error reading input image: {e}", file=sys.stderr)
        sys.exit(1)

    # Check if image is likely sRGB (8-bit or 16-bit integer types are usually sRGB)
    is_srgb = image_raw.dtype in (np.uint8, np.uint16)
    if is_srgb:
        print("Input image appears to be sRGB. Linearizing...")
        # Convert to float in [0,1]
        if image_raw.dtype == np.uint8:
            image_float = image_raw.astype(np.float64) / 255.0
        elif image_raw.dtype == np.uint16:
            image_float = image_raw.astype(np.float64) / 65535.0
        else:
            image_float = image_raw.astype(np.float64)

        # sRGB to linear conversion
        def srgb_to_linear(c):
            c = np.clip(c, 0.0, 1.0)
            return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)

        image_raw = srgb_to_linear(image_float)
    else:
        print("Input image is assumed to be linear RGB.")

    # --- Prepare Image Data ---
    # Convert to float64 for precision, handle different input types
    if image_raw.dtype == np.uint8:
        image_linear = image_raw.astype(np.float64) / 255.0
    elif image_raw.dtype == np.uint16:
        image_linear = image_raw.astype(np.float64) / 65535.0
    elif image_raw.dtype == np.float32:
        image_linear = image_raw.astype(np.float64)
    elif image_raw.dtype == np.float64:
        image_linear = image_raw
    else:
        print(f"Error: Unsupported image data type: {image_raw.dtype}", file=sys.stderr)
        sys.exit(1)

    # Discard alpha channel if present
    if image_linear.shape[-1] == 4:
        print("Discarding alpha channel.")
        image_linear = image_linear[..., :3]
    elif image_linear.shape[-1] != 3:
        print(
            f"Error: Input image must be RGB (shape {image_linear.shape} not supported).",
            file=sys.stderr,
        )
        sys.exit(1)

    # Ensure input is non-negative
    image_linear = np.maximum(image_linear, 0.0)
    print(f"Input image dimensions: {image_linear.shape}")
    image_width_px = image_linear.shape[1]

    # --- Pipeline Steps ---

    # 1. Convert Linear RGB to Log Exposure (LogE)
    # Map linear 0.18 to the specified middle_gray_logE from the datasheet
    print("Converting linear RGB to Log Exposure...")
    middle_gray_logE = float(datasheet.properties.calibration.middle_gray_logE)
    # Add epsilon inside log10 to handle pure black pixels
    log_exposure_rgb = middle_gray_logE + np.log10(image_linear / 0.18 + EPSILON)
    # Note: Values below 0.18 * 10**(hd_data['LogE'][0] - middle_gray_logE)
    # or above 0.18 * 10**(hd_data['LogE'][-1] - middle_gray_logE)
    # will map outside the H&D curve's defined LogE range and rely on clamping/extrapolation.

    # 2. Apply H&D Curves (Tonal Mapping + Balance Shifts + Gamma/Contrast)
    print("Applying H&D curves...")
    density_rgb = apply_hd_curves(
        log_exposure_rgb,
        datasheet.processing,
        datasheet.properties.curves.hd,
        middle_gray_logE,
    )

    # 3. Convert Density back to Linear Transmittance
    # Higher density means lower transmittance
    print("Converting density to linear transmittance...")
    # Add density epsilon? Usually density floor (Dmin) handles this.
    linear_transmittance = 10.0 ** (-density_rgb)
    # Normalize transmittance? Optional, assumes Dmin corresponds roughly to max transmittance 1.0
    # Could normalize relative to Dmin from the curves if needed.
    # linear_transmittance = linear_transmittance / (10.0**(-np.array([hd_data['Density_R'][0], hd_data['Density_G'][0], hd_data['Density_B'][0]])))
    linear_transmittance = np.clip(linear_transmittance, 0.0, 1.0)

    # 4. Apply Spatial Effects (Diffusion Blur, Halation)
    print("Applying spatial effects (diffusion, halation)...")
    # Apply these effects in the linear domain
    linear_post_spatial = apply_spatial_effects(
        linear_transmittance,
        datasheet.info.format_mm,
        datasheet.properties.couplers,
        datasheet.properties.interlayer,
        datasheet.properties.halation,
        image_width_px,
    )

    # 5. Apply Saturation Adjustment (Approximating Coupler Effects)
    print("Applying saturation adjustment...")
    coupler_amount = datasheet.properties.couplers.amount
    # Assuming coupler_amount directly scales saturation factor.
    # Values > 1 increase saturation, < 1 decrease.
    linear_post_saturation = apply_saturation_rgb(linear_post_spatial, coupler_amount)

    # --- Final Output Conversion ---
    print("Converting to output format...")
    # Clip final result and convert to uint8
    if args.output_image.lower().endswith(".tiff"):
        output_image_uint8 = (
            np.clip(linear_post_saturation, 0.0, 1.0) * 65535.0
        ).astype(np.uint16)
    else:
        output_image_uint8 = (np.clip(linear_post_saturation, 0.0, 1.0) * 255.0).astype(
            np.uint8
        )

    # --- Save Output Image ---
    print(f"Saving output image: {args.output_image}")
    try:
        if args.output_image.lower().endswith((".tiff", ".tif")):
            # Use imageio for standard formats
            iio.imwrite(args.output_image, output_image_uint8)
        elif args.output_image.lower().endswith(".png"):
            iio.imwrite(args.output_image, output_image_uint8, format="PNG")
        else:
            iio.imwrite(args.output_image, output_image_uint8, quality=95)
        print("Done.")
    except Exception as e:
        print(f"Error writing output image: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()