#!/usr/bin/env python3
"""Export a spektrafilm release into a data pack for the native darktable module.

This is the version-upgrade mechanism: when a new spektrafilm release comes out,

    pip install <new spektrafilm>          (or pip install -e <checkout>)
    python spektrafilm_export_data.py -o ~/.config/darktable/spektrafilm

and the darktable module picks up the new profiles / data on restart. The
module itself contains only the *algorithms* (which track the spektrafilm
model version recorded in pack.json); all measured data lives in this pack.

Pack layout:
    pack.json           model constants, CMFS, spectral locus, illuminant SPDs,
                        dichroic filter curves, neutral print filter database
    spectra_lut.f32     hanatos2025 irradiance spectra LUT (header + float32)
    profiles/*.json     verbatim spektrafilm stock profiles (CC BY-SA 4.0)
"""

import argparse
import json
import math
import shutil
import struct
import sys
from pathlib import Path

import numpy as np


def _chan3(v, bw):
    """Normalise a per-channel parameter to exactly three values.

    For a colour stock the three are passed through. For a single-emulsion (B&W)
    stock there is only one channel: upstream reaches every per-channel constant
    through match_channels(values, n_ch), which at n_ch == 1 returns values[:1] --
    the FIRST channel, for all of them. Since the C engine runs on the widened
    3-channel profile, that first value has to be replicated, or a B&W frame is
    rendered with chromatic parameters. Schema defaults make this concrete:
    rms_granularity is (6, 8, 10), scatter_core_um (2.2, 2.0, 1.6),
    halation_strength (0.05, 0.015, 0.0) -- none of which are achromatic.

    None-padded singles ([x, None, None]) collapse the same way.
    """
    v = [x for x in v if x is not None]
    if not v:
        return v
    if bw or len(v) == 1:
        return [v[0]] * 3
    return list(v)


def _grain_export(params, bw):
    gr = params.film_render.grain
    out = {
        "rms_granularity": _chan3(gr.rms_granularity, bw),
        "uniformity": _chan3(gr.uniformity, bw),
        "density_min": _chan3(gr.density_min, bw),
    }
    if hasattr(gr, "particle_scale_sublayers"):
        # sub-layer scales, not channels -- never collapsed
        out["particle_scale_sublayers"] = list(gr.particle_scale_sublayers)
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("-o", "--output", required=True, help="output pack directory")
    args = ap.parse_args()

    try:
        import spektrafilm  # noqa: F401
        from importlib.metadata import version as dist_version
        import importlib.resources as pkg_resources
        from spektrafilm.config import SPECTRAL_SHAPE, STANDARD_OBSERVER_CMFS, LOG_EXPOSURE
        from spektrafilm.model.illuminants import standard_illuminant
        from spektrafilm.model.color_filters import DichroicFilters
        from spektrafilm.utils.io import read_neutral_print_filters
        from spektrafilm.utils.spectral_upsampling import _load_hanatos2025_spectra_lut
        from spektrafilm.utils.gamut_compression import spectral_locus_xy
    except ImportError as err:
        print(f"error: spektrafilm must be importable ({err})", file=sys.stderr)
        return 1

    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    (out / "profiles").mkdir(exist_ok=True)

    wl = SPECTRAL_SHAPE.wavelengths
    version = dist_version("spektrafilm")

    # --- illuminants (normalized SPDs on the model wavelength grid) ---------
    illuminant_names = ["D50", "D55", "D65", "D75", "T", "TH-KG3", "TH-KG3-L", "K75P"]
    illuminants = {}
    for name in illuminant_names:
        try:
            illuminants[name] = np.asarray(standard_illuminant(name), dtype=float).tolist()
        except Exception as err:  # keep the pack usable if one SPD is missing
            print(f"warning: skipping illuminant {name}: {err}", file=sys.stderr)

    # --- enlarger dichroic filters ------------------------------------------
    filter_brands = ["thorlabs", "edmund_optics", "durst_digital_light", "custom"]
    dichroics = {}
    for brand in filter_brands:
        try:
            dichroics[brand] = np.asarray(DichroicFilters(brand=brand).filters, dtype=float).tolist()
        except Exception as err:
            print(f"warning: skipping dichroic brand {brand}: {err}", file=sys.stderr)

    # --- neutral print filter database --------------------------------------
    try:
        neutral_filters = read_neutral_print_filters()
    except FileNotFoundError:
        neutral_filters = {}

    # --- per-film digested render defaults -----------------------------------
    # Stock-specific tuning (DIR coupler gamma matrices, halation presets)
    # lives in spektrafilm's params_builder. Exporting the digested values per
    # film keeps that tuning in the data pack so the native module needs no
    # hardcoded per-stock tables and upgrades cleanly with new releases.
    from spektrafilm.runtime.params_builder import init_params, digest_params

    film_render_defaults = {}
    for res in pkg_resources.files("spektrafilm.data.profiles").iterdir():
        if not res.name.endswith(".json"):
            continue
        stock = res.name[:-5]
        with pkg_resources.as_file(res) as p:
            info = json.loads(Path(p).read_text()).get("info", {})
        if info.get("stage") != "filming":
            continue
        target_print = info.get("target_print") or "kodak_portra_endura"
        try:
            params = digest_params(init_params(film_profile=stock, print_profile=target_print))
        except Exception as err:
            print(f"warning: could not digest defaults for {stock}: {err}", file=sys.stderr)
            continue
        dc = params.film_render.dir_couplers
        ha = params.film_render.halation
        bw = bool(getattr(params.film, "is_bw", False))
        if bw:
            # Single emulsion: upstream's matrix is 1x1 self-inhibition
            # (compute_dir_couplers_matrix populates M_inter only for n_ch == 3),
            # so it uses gamma_samelayer_rgb[0] and its interlayer entries are
            # inert. They are NOT zero in couplers.toml -- defaults.bw.negative
            # carries the colour values verbatim, deliberately -- so on the
            # widened 3-channel profile they have to be zeroed here or every
            # channel picks up its whole matrix COLUMN instead of the diagonal.
            g0 = dc.gamma_samelayer_rgb[0]
            dc.gamma_samelayer_rgb = (g0, g0, g0)
            dc.gamma_interlayer_r_to_gb = (0.0, 0.0)
            dc.gamma_interlayer_g_to_rb = (0.0, 0.0)
            dc.gamma_interlayer_b_to_rg = (0.0, 0.0)
        film_render_defaults[stock] = {
            "dir_couplers": {
                "gamma_samelayer_rgb": list(dc.gamma_samelayer_rgb),
                "gamma_interlayer_r_to_gb": list(dc.gamma_interlayer_r_to_gb),
                "gamma_interlayer_g_to_rb": list(dc.gamma_interlayer_g_to_rb),
                "gamma_interlayer_b_to_rg": list(dc.gamma_interlayer_b_to_rg),
                "diffusion_size_um": dc.diffusion_size_um,
                "diffusion_tail_um": dc.diffusion_tail_um,
                "diffusion_tail_weight": dc.diffusion_tail_weight,
                # Langmuir saturating couplers (spektrafilm dev/0.4+); absent
                # on 0.3.x, in which case the engine uses the linear model
                **({"langmuir_donor_k_rgb": _chan3(dc.langmuir_donor_k_rgb, bw),
                    "langmuir_receiver_k_rgb": _chan3(dc.langmuir_receiver_k_rgb, bw)}
                   if hasattr(dc, "langmuir_donor_k_rgb") else {}),
            },
            "grain": _grain_export(params, bw),
            "halation": {
                "strength": _chan3(ha.halation_strength, bw),
                "first_sigma_um": _chan3(ha.halation_first_sigma_um, bw),
                "scatter_core_um": _chan3(ha.scatter_core_um, bw),
                "scatter_tail_um": _chan3(ha.scatter_tail_um, bw),
                "scatter_tail_weight": _chan3(ha.scatter_tail_weight, bw),
            },
        }

    pack = {
        "pack_format": 1,
        "spektrafilm_version": version,
        "film_render_defaults": film_render_defaults,
        "wavelengths": np.asarray(wl, dtype=float).tolist(),
        "log_exposure": np.asarray(LOG_EXPOSURE, dtype=float).tolist(),
        "cmfs": np.asarray(STANDARD_OBSERVER_CMFS[:], dtype=float).tolist(),
        "spectral_locus_xy": np.asarray(spectral_locus_xy(), dtype=float).tolist(),
        "illuminants": illuminants,
        "dichroic_filters": dichroics,
        "neutral_print_filters": neutral_filters,
    }
    def _sanitize(obj):
        """JSON has no NaN/Inf — emit null instead (readers treat null as NaN)."""
        if isinstance(obj, dict):
            return {k: _sanitize(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_sanitize(v) for v in obj]
        if isinstance(obj, float) and not math.isfinite(obj):
            return None
        return obj

    (out / "pack.json").write_text(json.dumps(_sanitize(pack)))
    print(f"wrote {out / 'pack.json'} (spektrafilm {version})")

    # --- spectral upsampling LUT ---------------------------------------------
    # Written as float16, which is what upstream computes, ships and uses --
    # src/spektrafilm/data/luts/spectral_upsampling/*.npy are float16 arrays.
    # The previous float32 conversion here doubled the file (11.4 MiB against
    # 5.7) for precision upstream never had.
    #
    # The header carries the LUT's identity as well as its shape. Upstream
    # revises this table often (arctic2026 alpha, alpha02, beta01..beta04 so
    # far) and each revision changes every render, so the engine has to be able
    # to tell one from another rather than silently rendering differently. The
    # content hash is authoritative -- the name is discovered when possible and
    # is only there to make the mismatch message readable.
    lut_raw = _load_hanatos2025_spectra_lut()
    lut = np.ascontiguousarray(lut_raw, dtype=np.float16)

    lut_name = "unknown"
    for attr in ("SPECTRAL_UPSAMPLING_LUT", "HANATOS2025_LUT_NAME"):
        try:
            from spektrafilm import config as _cfg
            lut_name = str(getattr(_cfg, attr))
            break
        except (ImportError, AttributeError):
            pass
    if lut_name == "unknown":
        # fall back to the newest LUT shipped in the package
        try:
            lutdir = pkg_resources.files("spektrafilm.data.luts.spectral_upsampling")
            names = sorted(r.name[:-4] for r in lutdir.iterdir() if r.name.endswith(".npy"))
            if names:
                lut_name = names[-1]
        except Exception:
            pass

    # FNV-1a over the raw bytes; the engine reads this value, it does not
    # recompute it, so any stable hash will do
    h = 2166136261
    for b in lut.tobytes():
        h = ((h ^ b) * 16777619) & 0xFFFFFFFF
    lut_id = f"{lut_name}@{version}".encode("utf-8")[:255]

    with open(out / "spectra_lut.f32", "wb") as fh:
        fh.write(b"SFS2")
        fh.write(struct.pack("<i", 2))                 # header version
        fh.write(struct.pack("<iii", *lut.shape))
        fh.write(struct.pack("<i", 1))                 # 0 = float32, 1 = float16
        fh.write(struct.pack("<I", h))                 # content hash
        fh.write(struct.pack("<i", len(lut_id)))
        fh.write(lut_id)
        fh.write(lut.tobytes())
    print(f"wrote {out / 'spectra_lut.f32'} shape {lut.shape} float16 "
          f"id={lut_id.decode()} hash={h:08x}")

    # --- stock profiles ------------------------------------------------------
    # Colour profiles are copied verbatim. Single-emulsion B&W profiles
    # (channel_model == "bw", spektrafilm dev/0.4+) are widened to the 3-channel
    # layout the C engine expects: per-channel arrays are triplicated and
    # channel_density is divided by 3, so the spectral sum over channels equals
    # the single emulsion's density spectrum. info.channel_model stays "bw" so
    # the module can couple the grain across channels and collapse per-channel
    # constants.
    #
    # Development-time families are NOT collapsed here any more. Some B&W stocks
    # carry one density curve, base+fog spectrum and curve-model row per
    # development time (kodak_doublex: 4/5/6.5/9/12 min, kodak_2302: 2/3.5/5/7/9
    # min). Collapsing to upstream's default middle member at export time threw
    # that away and made the choice unreachable from the module; keeping the
    # family lets the module do what select_development_time() does, at render
    # time, from a slider. `development_time` carries the full list of times, and
    # its length is what tells the loader whether the arrays are a family or an
    # already-widened single member -- so old packs keep loading unchanged.
    #
    # A family is exported with the development axis LAST on the 2-D arrays
    # (density_curves (n_le, n_dev), base_density (n_wl, n_dev)) and FIRST on the
    # curve model (centers (n_dev, n_layers)), which is how upstream stores them;
    # the module selects a member and widens to 3 channels itself.
    profile_dir = pkg_resources.files("spektrafilm.data.profiles")
    n = nbw = nfam = 0
    for res in profile_dir.iterdir():
        if not res.name.endswith(".json"):
            continue
        with pkg_resources.as_file(res) as p:
            prof = json.loads(p.read_text())
        if prof.get("info", {}).get("channel_model") == "bw":
            d = prof["data"]
            times = d.get("development_time") or []
            n_dev = len(times)
            curves = d.get("density_curves") or []
            family = (n_dev > 1 and curves and isinstance(curves[0], list)
                      and len(curves[0]) == n_dev)
            if family:
                # Leave density_curves, base_density and the curve model as the
                # full family; only sanity-check that the model row count agrees,
                # since a mismatch there is what silently summed N development
                # fits together as if they were N sub-layers of one curve.
                model = d.get("density_curves_model") or {}
                for key in ("centers", "amplitudes", "sigmas", "alphas"):
                    arr = model.get(key)
                    if arr and len(arr) != n_dev:
                        print(f"warning: {res.name}: density_curves_model.{key} has "
                              f"{len(arr)} rows but development_time has {n_dev}; "
                              f"not a development family, leaving as-is", file=sys.stderr)
                        family = False
                        break
            if family:
                nfam += 1
            else:
                # Single member (or no family at all): widen to 3 channels. The
                # curve model's one row is REPEATED across 3 identical outer rows,
                # not tripled element-wise, so each broadcast channel sums the
                # same complete set of terms.
                model = d.get("density_curves_model")
                if model:
                    for key in ("centers", "amplitudes", "sigmas", "alphas"):
                        arr = model.get(key)
                        if arr and len(arr) == 1:
                            model[key] = arr * 3
                for key in ("log_sensitivity", "density_curves"):
                    if key in d and d[key] and isinstance(d[key][0], list) \
                       and len(d[key][0]) == 1:
                        d[key] = [row * 3 for row in d[key]]
            # log_sensitivity is one panchromatic curve either way
            if "log_sensitivity" in d and d["log_sensitivity"] \
               and isinstance(d["log_sensitivity"][0], list) and len(d["log_sensitivity"][0]) == 1:
                d["log_sensitivity"] = [row * 3 for row in d["log_sensitivity"]]
            if "channel_density" in d and d["channel_density"] \
               and isinstance(d["channel_density"][0], list) and len(d["channel_density"][0]) == 1:
                d["channel_density"] = [
                    [None if row[0] is None else row[0] / 3.0] * 3
                    for row in d["channel_density"]
                ]
            nbw += 1
        (out / "profiles" / res.name).write_text(json.dumps(prof))
        n += 1
    print(f"copied {n} profiles ({nbw} B&W, {nfam} carrying a development-time family)")
    return 0


if __name__ == "__main__":
    sys.exit(main())