"""
CZI -> TIFF converter.

Reads a Carl Zeiss .czi microscopy file and writes one TIFF per
(scene, time, channel, z) plane. Designed to run both as a CLI script and
inside a Hugging Face Space (Gradio app at the bottom, guarded by __main__).

Why TIFF (and why files can look "blank"):
Microscopy data is usually 16-bit, but real signal often occupies only a small
fraction of the 0..65535 range (e.g. peaks around a few hundred). If such a
file is written with raw values and opened in an ordinary viewer (Preview,
Windows Photos, browsers, ZIP thumbnails), the viewer maps 0..65535 onto
0..255, so a peak of ~800 becomes pixel value ~3: visually indistinguishable
from black. The file is NOT empty; it just renders as black.

To avoid that, this converter defaults to writing VIEWABLE TIFFs: each plane is
linearly stretched to the full 16-bit range so it displays correctly in any
viewer. The stretch is recorded so the transform is transparent. For
quantitative work where absolute pixel values must be preserved bit-for-bit,
pass viewable=False to get RAW output (native bit depth, no scaling); those
files additionally carry ImageJ display-range tags so scientific tools
(Fiji/ImageJ, napari, QuPath) auto-scale them on open.

Dependencies:
    pip install czifile imagecodecs numpy tifffile pillow
(Gradio app also needs: pip install gradio)
"""

import os
import re
import numpy as np
import tifffile
import czifile


# ----------------------------------------------------------------------------
# Core conversion
# ----------------------------------------------------------------------------
def _canonical_array(path):
    """
    Return (array, axes_string) in the CZI canonical layout, NOT squeezed.
    czifile >= 2026 removed the public .axes/.shape attributes, so we read the
    full (unsqueezed) array which always follows the fixed canonical order
    'SCTZYX0' (Scene, Channel, Time, Z, Y, X, Samples). We then map by name.
    """
    czi = czifile.CziFile(path)
    try:
        # Force the full, predictable layout instead of relying on squeeze.
        czi._squeeze = False
        arr = czi.asarray()
        # Read the true axis order if exposed; otherwise fall back to the
        # observed canonical layout for this czifile version.
        axes = getattr(czi, "axes", None)
    finally:
        czi.close()

    # czifile's full (unsqueezed) array follows a fixed dimension order. For the
    # current library version this is S, C, T, Z, Y, X, Samples. The trailing
    # axis is samples-per-pixel (1=gray, 3=RGB, 4=RGBA).
    if not axes or len(axes) != arr.ndim:
        axes = "SCTZYX0"
    if arr.ndim != len(axes):
        # Fallback: pad/trim by leading singleton dims so the named map still works.
        while arr.ndim < len(axes):
            arr = arr[np.newaxis, ...]
        if arr.ndim > len(axes):
            # Collapse any extra leading singleton dims.
            extra = arr.ndim - len(axes)
            for _ in range(extra):
                if arr.shape[0] == 1:
                    arr = arr[0]
                else:
                    break
    return arr, axes


def _to_raw_dtype(plane):
    """
    Return the plane in a TIFF-writable dtype, PRESERVING NATIVE BIT DEPTH and
    ABSOLUTE INTENSITY. No contrast stretching, no normalization, no 8-bit
    down-conversion. uint8/uint16/uint32/float32 pass through unchanged. Other
    integer types are cast to the smallest standard container that holds them
    without altering values. float64 is preserved as float32.
    """
    plane = np.asarray(plane)
    dt = plane.dtype

    if dt in (np.uint8, np.uint16, np.uint32, np.float32):
        return plane
    if dt == np.float64:
        return plane.astype(np.float32)

    if np.issubdtype(dt, np.integer):
        info = np.iinfo(dt)
        if info.min < 0:
            if info.bits <= 16:
                return plane.astype(np.int16)
            return plane.astype(np.int32)
        if info.bits <= 8:
            return plane.astype(np.uint8)
        if info.bits <= 16:
            return plane.astype(np.uint16)
        return plane.astype(np.uint32)

    return plane.astype(np.float32)


def _stretch_to_viewable(plane):
    """
    Linearly stretch a single-sample plane to the full unsigned range of an
    8- or 16-bit container so it displays correctly in ANY image viewer.
    The output bit depth matches the input's natural container (8-bit stays
    8-bit, everything else becomes 16-bit) so contrast is maximized without
    needlessly inflating file size. Returns (stretched_array, lo, hi) where
    lo/hi are the original min/max that were mapped to black/white.
    """
    a = np.asarray(plane).astype(np.float64)
    lo = float(a.min())
    hi = float(a.max())

    # Choose an 8-bit container only for genuine 8-bit input; otherwise 16-bit.
    if np.asarray(plane).dtype == np.uint8:
        out_max, out_dtype = 255.0, np.uint8
    else:
        out_max, out_dtype = 65535.0, np.uint16

    if hi > lo:
        a = (a - lo) * (out_max / (hi - lo))
    else:
        # Flat plane: leave at zero (truly empty) rather than divide by zero.
        a = np.zeros_like(a)
    return a.round().clip(0, out_max).astype(out_dtype), lo, hi


def _channel_names(path, n_channels):
    """Best-effort channel names from metadata; falls back to C0, C1, ..."""
    names = []
    try:
        czi = czifile.CziFile(path)
        try:
            md = czi.metadata()
        finally:
            czi.close()
        # Names can repeat in metadata; keep first occurrence order, deduped.
        found = re.findall(r'<Channel[^>]*?Name="([^"]+)"', md or "")
        seen = []
        for f in found:
            if f not in seen:
                seen.append(f)
        names = seen
    except Exception:
        names = []

    out = []
    for i in range(n_channels):
        if i < len(names):
            safe = re.sub(r"[^A-Za-z0-9._-]+", "_", names[i]).strip("_")
            out.append(safe or f"C{i}")
        else:
            out.append(f"C{i}")
    return out


def _write_plane(fpath, plane, photometric, viewable):
    """
    Write a single plane to TIFF.
    viewable=True  : contrast-stretch to full container range so the file is
                     visible in any viewer. Stored as uint8/uint16.
    viewable=False : write raw values (native bit depth, absolute intensity
                     preserved) and embed ImageJ display-range tags so
                     scientific viewers auto-scale on open.
    """
    if photometric == "rgb":
        # RGB/RGBA: write as-is. Color data already spans a sensible range and
        # per-channel stretching would shift colors, so we never stretch it.
        img = _to_raw_dtype(plane)
        tifffile.imwrite(fpath, img, photometric="rgb", compression="deflate")
        return

    if viewable:
        img, lo, hi = _stretch_to_viewable(plane)
        # Record the original intensity window we mapped from, for traceability.
        tifffile.imwrite(
            fpath,
            img,
            photometric="minisblack",
            compression="deflate",
            metadata={"original_min": lo, "original_max": hi, "stretched": True},
        )
    else:
        img = _to_raw_dtype(plane)
        lo, hi = float(np.asarray(img).min()), float(np.asarray(img).max())
        # ImageJ tags let Fiji/napari/QuPath auto-scale without altering pixels.
        # ImageJ format requires uint8/uint16/float32; cast uint32 safely up.
        if img.dtype == np.uint32:
            img = img.astype(np.float32)
        tifffile.imwrite(
            fpath,
            img,
            imagej=True,
            metadata={"mode": "grayscale", "min": lo, "max": hi},
        )


def convert_czi_to_tiff(input_path, output_dir, viewable=True):
    """
    Convert one .czi file to a set of TIFF files.
    viewable=True (default): files are contrast-stretched so they display
        correctly in any viewer (fixes the "all files look blank" problem
        caused by low-intensity 16-bit microscopy data).
    viewable=False: raw absolute-intensity output for quantitative analysis,
        with ImageJ display-range tags for scientific viewers.
    Returns a list of written TIFF file paths.
    """
    os.makedirs(output_dir, exist_ok=True)
    arr, axes = _canonical_array(input_path)

    idx = {ax: axes.index(ax) for ax in axes}
    nS = arr.shape[idx["S"]]
    nT = arr.shape[idx["T"]]
    nC = arr.shape[idx["C"]]
    nZ = arr.shape[idx["Z"]]
    n_samples = arr.shape[idx["0"]]  # samples per pixel (1=gray, 3=RGB, 4=RGBA)

    ch_names = _channel_names(input_path, nC)
    base = os.path.splitext(os.path.basename(input_path))[0]
    written = []

    for s in range(nS):
        for t in range(nT):
            for c in range(nC):
                for z in range(nZ):
                    # Slice down to a single plane (Y, X, samples).
                    sl = [slice(None)] * arr.ndim
                    sl[idx["S"]] = s
                    sl[idx["T"]] = t
                    sl[idx["C"]] = c
                    sl[idx["Z"]] = z
                    plane = arr[tuple(sl)]  # leaves Y, X, and samples axes

                    plane = np.squeeze(plane)
                    if plane.ndim == 3 and plane.shape[-1] in (3, 4):
                        photometric = "rgb"
                        out_plane = plane
                    elif plane.ndim == 2:
                        photometric = "minisblack"
                        out_plane = plane
                    else:
                        # Unexpected extra dims: collapse to 2D defensively.
                        photometric = "minisblack"
                        out_plane = plane.reshape(plane.shape[-2], plane.shape[-1])

                    parts = [base]
                    if nS > 1:
                        parts.append(f"S{s}")
                    if nT > 1:
                        parts.append(f"T{t}")
                    parts.append(ch_names[c])
                    if nZ > 1:
                        parts.append(f"Z{z:02d}")
                    fname = "_".join(parts) + ".tiff"
                    fpath = os.path.join(output_dir, fname)

                    _write_plane(fpath, out_plane, photometric, viewable)
                    written.append(fpath)

    return written


# ----------------------------------------------------------------------------
# Hugging Face Gradio app
# ----------------------------------------------------------------------------
def _build_gradio_app():
    import tempfile
    import zipfile
    import gradio as gr

    def _to_preview_uint8(img):
        """
        8-bit grayscale preview for on-screen display only. Independent of the
        saved files: it always stretches so the browser shows real content.
        """
        a = np.asarray(img)
        if a.ndim == 3 and a.shape[-1] in (3, 4):
            a = a[..., :3].astype(np.float64)
            chans = []
            for k in range(a.shape[-1]):
                ch = a[..., k]
                lo, hi = float(ch.min()), float(ch.max())
                ch = (ch - lo) * (255.0 / (hi - lo)) if hi > lo else np.zeros_like(ch)
                chans.append(ch.clip(0, 255).astype(np.uint8))
            return np.stack(chans, axis=-1)
        a = a.astype(np.float64)
        lo, hi = float(a.min()), float(a.max())
        a = (a - lo) * (255.0 / (hi - lo)) if hi > lo else np.zeros_like(a)
        return a.clip(0, 255).astype(np.uint8)

    def _process(file_obj):
        if file_obj is None:
            return None, [], "Please upload a .czi file."
        tmp = tempfile.mkdtemp()
        out_dir = os.path.join(tmp, "tiff")
        # Raw output is now the default and only mode.
        tiffs = convert_czi_to_tiff(file_obj.name, out_dir, viewable=False)

        zip_path = os.path.join(tmp, "czi_tiffs.zip")
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for p in tiffs:
                zf.write(p, arcname=os.path.basename(p))

        gallery = []
        for p in tiffs:
            arr = tifffile.imread(p)
            gallery.append((_to_preview_uint8(arr), os.path.basename(p)))

        msg = (f"Converted {len(tiffs)} plane(s) to RAW TIFF "
               f"(absolute intensity preserved; ImageJ display tags embedded). "
               f"These may look black in generic viewers but open correctly in Fiji/napari.")
        return zip_path, gallery, msg

    grayscale_css = """
    .gradio-container { background: #ffffff !important; color: #000000 !important; }
    .gradio-container *, .prose, .prose * { color: #000000 !important; }
    .gradio-container h1, .gradio-container h2, .gradio-container h3,
    .gradio-container p, .gradio-container span, .gradio-container label,
    .gradio-container .markdown, .gradio-container .markdown * {
        color: #000000 !important;
    }
    .gallery img, .grid-wrap img { filter: grayscale(100%); background: #000000; }
    button.primary, .primary button, #convert-btn, #convert-btn button {
        background: #000000 !important;
        background-image: none !important;
        color: #ffffff !important;
        border-color: #000000 !important;
    }
    #convert-btn *, .primary button * { color: #ffffff !important; }
    """

    with gr.Blocks(title="CZI to TIFF Converter", css=grayscale_css,
                   theme=gr.themes.Default()) as demo:
        gr.Markdown(
            "# CZI to TIFF Converter\n"
            "Upload a `.czi` file to get a ZIP of TIFFs (one per scene/time/"
            "channel/Z plane).\n\n"
            "Files are written as **raw** TIFFs that preserve absolute pixel "
            "intensities for quantitative analysis (they carry ImageJ display "
            "tags for Fiji/napari/QuPath). Real 16-bit microscopy signal is "
            "often very faint, so these raw files can render as solid black in "
            "generic viewers even though the pixel data is intact; the preview "
            "below is contrast-stretched for on-screen display only."
        )
        inp = gr.File(label="CZI file", file_types=[".czi"])
        btn = gr.Button("Convert", variant="primary", elem_id="convert-btn")
        out_msg = gr.Textbox(label="Status")
        out_gallery = gr.Gallery(label="Preview (8-bit, display only)", columns=4, height="auto")
        out_file = gr.File(label="Download TIFFs (ZIP)")
        btn.click(_process, inputs=[inp], outputs=[out_file, out_gallery, out_msg])

    return demo


if __name__ == "__main__":
    import sys

    # CLI mode:  python czi_to_tiff.py input.czi [output_dir] [--viewable]
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    viewable_flag = "--viewable" in sys.argv
    if args and args[0].lower().endswith(".czi"):
        in_path = args[0]
        out_dir = args[1] if len(args) >= 2 else "tiff_output"
        files = convert_czi_to_tiff(in_path, out_dir, viewable=viewable_flag)
        mode = "viewable" if viewable_flag else "raw"
        print(f"Wrote {len(files)} TIFF file(s) ({mode}) to {out_dir}/")
        for f in files:
            print("  ", os.path.basename(f))
    else:
        # App mode (Hugging Face Space)
        _build_gradio_app().launch()