"""
CZI -> PNG converter.

Reads a Carl Zeiss .czi microscopy file and writes one PNG 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__).

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

import os
import re
import numpy as np
from PIL import Image
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
    'STCZYX0' (Scene, Time, Channel, 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_uint8(plane):
    """
    Convert a plane to uint8 PNG values while PRESERVING ABSOLUTE INTENSITY.

    No contrast stretching. uint8 data passes through unchanged. Higher bit
    depths (uint16, etc.) are rescaled only by the ratio of their full storage
    range to 255 (e.g. uint16 -> divide by 257), which preserves absolute
    intensity meaning rather than per-plane min/max normalization.
    """
    plane = np.asarray(plane)
    if plane.dtype == np.uint8:
        return plane

    if np.issubdtype(plane.dtype, np.integer):
        info = np.iinfo(plane.dtype)
        max_val = info.max  # full storage range, NOT this plane's observed max
        scaled = plane.astype(np.float64) * (255.0 / max_val)
        return scaled.round().clip(0, 255).astype(np.uint8)

    # Floating point: assume already in 0..1 if max <= 1, else 0..255.
    plane = plane.astype(np.float64)
    if plane.max() <= 1.0:
        plane = plane * 255.0
    return plane.round().clip(0, 255).astype(np.uint8)


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 convert_czi_to_png(input_path, output_dir):
    """
    Convert one .czi file to a set of PNG files.

    Returns a list of written PNG 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):
                        img8 = np.stack(
                            [_to_uint8(plane[..., k]) for k in range(plane.shape[-1])],
                            axis=-1,
                        )
                        mode = "RGB" if img8.shape[-1] == 3 else "RGBA"
                    elif plane.ndim == 2:
                        img8 = _to_uint8(plane)
                        mode = "L"
                    else:
                        # Unexpected extra dims: collapse to 2D defensively.
                        img8 = _to_uint8(plane.reshape(plane.shape[-2], plane.shape[-1]))
                        mode = "L"

                    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) + ".png"
                    fpath = os.path.join(output_dir, fname)

                    Image.fromarray(img8, mode=mode).save(fpath)
                    written.append(fpath)

    return written


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

    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, "png")
        pngs = convert_czi_to_png(file_obj.name, out_dir)

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

        # Gallery items: (grayscale image path, caption). Images are saved as
        # mode "L", so they display in pure black and white.
        gallery = [(p, os.path.basename(p)) for p in pngs]
        msg = f"Converted {len(pngs)} plane(s) to PNG."
        return zip_path, gallery, msg

    # Force a plain white/black look so previews read as grayscale microscopy.
    grayscale_css = """
    .gradio-container { background: #ffffff; color: #000000; }
    .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;
    }
    """

    with gr.Blocks(title="CZI to PNG Converter", css=grayscale_css) as demo:
        gr.Markdown("# CZI to PNG Converter\nUpload a `.czi` file to get a ZIP of black-and-white PNGs (one per channel/Z plane), with absolute pixel intensities preserved.")
        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 (black & white)", columns=4, height="auto")
        out_file = gr.File(label="Download PNGs (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_png.py input.czi [output_dir]
    if len(sys.argv) >= 2 and sys.argv[1].lower().endswith(".czi"):
        in_path = sys.argv[1]
        out_dir = sys.argv[2] if len(sys.argv) >= 3 else "png_output"
        files = convert_czi_to_png(in_path, out_dir)
        print(f"Wrote {len(files)} PNG file(s) to {out_dir}/")
        for f in files:
            print("  ", os.path.basename(f))
    else:
        # App mode (Hugging Face Space)
        _build_gradio_app().launch()