"""Offline medical-image filter comparisons from local volumes and DICOM."""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import json
import re
import textwrap
import sys

sys.path.insert(0, str(Path(__file__).resolve().parent))
from medical_io import CASE_NAMES, bundled_cases, load_medical, representative_plane

from image_filter import (
    CannyFilter, GaussianFilter, MedianFilter, SobelFilter, LaplacianFilter,
    GaborFilter, LawsFilter, LoGFilter, MeanFilter, WaveletFilter, Pipeline,
)

DEFAULT_DIR = Path(__file__).resolve().parents[1] / "example-output"
STYLE = {"font.family": "DejaVu Sans", "font.size": 11,
         "text.color": "#18283b", "axes.titlecolor": "#18283b",
         "figure.facecolor": "#f4f6f9", "savefig.facecolor": "#f4f6f9"}


def render_panel(name: str, image: np.ndarray, info: dict, output: Path) -> None:
    operations = [
        ("Mean 5×5", MeanFilter(5)),
        ("Gaussian σ=1.5", GaussianFilter(1.5)),
        ("Median 5×5", MedianFilter(5)),
        ("Sobel magnitude", SobelFilter()),
        ("Laplacian", LaplacianFilter()),
        ("Canny σ=1.2 (quantiles .1/.3)", CannyFilter(1.2, .1, .3, True)),
        ("Laws L5E5 energy r=3", LawsFilter("L5E5", energy_radius=3)),
        ("LoG σ=1.5", LoGFilter(1.5)),
        ("Gabor magnitude", GaborFilter(0.15)),
        ("SWT LH (db2)", WaveletFilter("LH", "db2")),
    ]
    # One-step composites preserve the existing independent comparisons.
    panels = [("Input", image)] + [
        (title, Pipeline([operation]).apply(image).image)
        for title, operation in operations
    ]
    with plt.rc_context(STYLE):
        figure, axes = plt.subplots(3, 4, figsize=(16, 12))
        figure.subplots_adjust(left=.025, right=.975, bottom=.035, top=.88,
                               hspace=.23, wspace=.10)
        figure.text(.025, .965, "IMAGE FILTER", fontsize=12, weight="bold", color="#426080")
        figure.text(.025, .925, f"{info['modality']} · {info['format']} · {name}", fontsize=23, weight="bold")
        for index, (axis, (title, values)) in enumerate(zip(axes.flat, panels)):
            limits = (0, 1) if values.dtype == np.bool_ else np.percentile(values, (1, 99))
            axis.imshow(values, cmap="gray", vmin=limits[0], vmax=limits[1])
            axis.set_title(f"{index:02d}  {title}", fontsize=12, loc="left", pad=9)
            axis.axis("off")
        note = axes.flat[-1]
        note.axis("off")
        selection = f"Axis {info['slice_axis']} / slice {info['slice_index']}" if info['slice_index'] is not None else "Native 2D frame"
        source_label = textwrap.fill(info['source'], width=30)
        note.text(.02, .95, "CASE / PROCESSING", transform=note.transAxes, weight="bold", fontsize=12)
        note.text(.02, .83,
            f"Source: {source_label}\n{selection} · {info['files']} file(s)\nInput: {image.shape[1]} × {image.shape[0]}\n{info['orientation']}\n\n1–99% window → [0, 1] input.\nAll filters act on this 2D slice.\nIndependent output display contrast.\nCanny: boolean edges.",
            transform=note.transAxes, fontsize=10, va="top", linespacing=1.55)
        figure.savefig(output, dpi=150, metadata={"Title": f"Image Filter comparison: {name}"})
        plt.close(figure)


def render_overview(originals: dict[str, tuple[np.ndarray, dict]], output: Path) -> None:
    rows = (len(originals) + 2) // 3
    with plt.rc_context(STYLE):
        figure, axes = plt.subplots(rows, 3, figsize=(12, 4 * rows + 1), squeeze=False)
        figure.subplots_adjust(left=.035, right=.965, bottom=.045, top=1 - 1.25 / (4 * rows + 1),
                               hspace=.25, wspace=.12)
        figure.suptitle("Image Filter · medical imaging examples", fontsize=20, weight="bold", y=.975)
        for axis in axes.flat:
            axis.axis("off")
        for axis, (name, (original, info)) in zip(axes.flat, originals.items()):
            axis.imshow(original, cmap="gray", vmin=0, vmax=1)
            axis.set_anchor("N")
            axis.set_title(f"{info['modality']} · {info['format']}\n{name}-filters.png", fontsize=12, loc="left")
        figure.savefig(output, dpi=150)
        plt.close(figure)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    selection = parser.add_mutually_exclusive_group()
    selection.add_argument("--all", action="store_true", help="All available local medical fixtures")
    selection.add_argument("--images", nargs="+", choices=CASE_NAMES, help="Select local medical cases")
    parser.add_argument("--input", type=Path, nargs="+", default=[], help="User NIfTI, NRRD or single DICOM paths")
    parser.add_argument("--dicom-series", type=Path, nargs="+", default=[], help="DICOM series directories (one volume each)")
    parser.add_argument("--series-uid", help="Select a series when a directory contains multiple series")
    parser.add_argument("--modality", default="Unknown", help="Modality label for user NIfTI/NRRD files")
    parser.add_argument("--slice-axis", type=int, choices=(0, 1, 2), default=2, help="Canonical voxel axis for volumes")
    parser.add_argument("--slice-index", type=int, help="Override deterministic non-empty slice selection")
    destination = parser.add_mutually_exclusive_group()
    destination.add_argument("--output", type=Path, help="Single PNG output")
    destination.add_argument("--output-dir", type=Path, default=None, help="Default: package example-output directory")
    parser.add_argument("--contact-sheet", action="store_true", help="Write source-overview.png")
    args = parser.parse_args(argv)
    args.names = list(CASE_NAMES) if args.all else list(dict.fromkeys(args.images or ([] if args.input or args.dicom_series else ['mri'])))
    count = len(args.names) + len(args.input) + len(args.dicom_series)
    if args.output and (count != 1 or args.output.suffix.lower() != '.png'):
        parser.error('--output requires exactly one source and a .png extension')
    if args.output and args.output.name == 'source-overview.png':
        parser.error('--output must not use the reserved overview filename')
    return args


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    catalog = bundled_cases()
    selections = []
    for name in args.names:
        if name not in catalog or not catalog[name][0].exists():
            if args.all:
                print(f"Unavailable bundled case: {name}; no download attempted")
                continue
            raise SystemExit(f"Unavailable bundled case: {name}")
        path, modality = catalog[name]
        selections.append((name, path, modality))
    for path in args.input + args.dicom_series:
        name = re.sub(r'[^a-z0-9_-]+', '-', path.name.lower().removesuffix('.gz').removesuffix('.nii').removesuffix('.nrrd').removesuffix('.dcm')).strip('-')
        selections.append((name or 'case', path, args.modality))
    if len({name for name, _, _ in selections}) != len(selections):
        raise SystemExit('Source basenames collide; run cases separately with distinct --output names')
    if not selections:
        raise SystemExit('No local medical fixtures available')
    directory = args.output.parent if args.output else (args.output_dir or DEFAULT_DIR)
    directory.mkdir(parents=True, exist_ok=True)
    results = {}
    for name, path, modality in selections:
        try:
            case = load_medical(path, modality, args.series_uid)
            image, info = representative_plane(case, args.slice_axis, args.slice_index)
        except (ValueError, FileNotFoundError, RuntimeError, ImportError) as exc:
            raise SystemExit(f"{name}: {exc}") from exc
        output = args.output or directory / f'{name}-filters.png'
        render_panel(name, image, info, output)
        results[name] = (image, info)
        print(output.resolve())
    (directory / 'cases.json').write_text(json.dumps({name: info for name, (_, info) in results.items()}, indent=2) + '\n')
    if args.contact_sheet:
        output = directory / 'source-overview.png'
        render_overview(results, output)
        print(output.resolve())
    if args.all:
        print('Bundled NRRD coverage unavailable; use --input /path/to/case.nrrd. No downloads performed.')


if __name__ == "__main__":
    main()