import pathlib

import matplotlib
import matplotlib.pyplot as plt

import palom


def align_he(
    p1: str | pathlib.Path,
    p2: str | pathlib.Path,
    out_dir: str | pathlib.Path,
    out_name: str = None,
    thumbnail_channel1: int = 1,
    thumbnail_channel2: int = 2,
    channel1: int = 0,
    channel2: int = 2,
    px_size1: float = None,
    px_size2: float = None,
    n_keypoints: int = 10_000,
    auto_mask=True,
    only_coarse: bool = False,
    only_qc: bool = False,
    viz_coarse_napari: bool = False,
    multi_res: bool = False,
    multi_obj: bool = False,
    multi_obj_kwarg: dict = None,
):
    assert not (multi_res and multi_obj), (
        "setting both `multi_res` and `multi_obj` to `True` is not supported,"
        " choose at most one"
    )
    out_dir, p1, p2 = pathlib.Path(out_dir), pathlib.Path(p1), pathlib.Path(p2)
    if out_name is None:
        out_name = f"{p2.stem}-registered.ome.tif"
    out_path = out_dir / out_name
    assert "".join(out_path.suffixes[-2:]) in (".ome.tif", ".ome.tiff")
    out_path.parent.mkdir(exist_ok=True, parents=True)
    set_matplotlib_font(font_size=8)

    r1 = get_reader(p1)(p1, pixel_size=px_size1)
    r2 = get_reader(p2)(p2, pixel_size=px_size2)

    LEVEL = 0
    aligner = palom.align.get_aligner(
        r1,
        r2,
        level1=LEVEL,
        level2=LEVEL,
        channel1=channel1,
        channel2=channel2,
        # make thumbnail level pair based on pixel_size
        thumbnail_level1=None,
        thumbnail_channel1=thumbnail_channel1,
        thumbnail_channel2=thumbnail_channel2,
    )

    aligner.coarse_register_affine(
        n_keypoints=n_keypoints,
        test_flip=True,
        test_intensity_invert=True,
        auto_mask=auto_mask,
    )
    fig, ax = plt.gcf(), plt.gca()
    fig.suptitle(f"{p2.name} (coarse alignment)", fontsize=8)
    ax.set_title(f"{p1.name} - {p2.name}", fontsize=6)
    im_h, im_w = ax.images[0].get_array().shape
    set_subplot_size(im_w / 288, im_h / 288, ax=ax)
    ax.set_anchor("N")
    # use 0.5 inch on the top for figure title
    fig.subplots_adjust(top=1 - 0.5 / fig.get_size_inches()[1])
    save_all_figs(out_dir=out_dir / "qc", format="jpg", dpi=144)

    if viz_coarse_napari:
        _ = viz_coarse(r1, r2, LEVEL, LEVEL, channel1, channel2, aligner.affine_matrix)

    if not only_coarse:
        # the default
        if not (multi_res or multi_obj):
            aligner.compute_shifts()

            fig = aligner.plot_shifts()
            fig.suptitle(f"{p2.name} (block shift distance)", fontsize=8)
            fig.axes[0].set_title(p1.name, fontsize=6)
            save_all_figs(out_dir=out_dir / "qc", format="png")

            aligner.constrain_shifts()
            block_mx = aligner.block_affine_matrices_da
        elif multi_res:
            mr_aligner = palom.align_multi_res.MultiResAligner(
                r1,
                r2,
                level1=LEVEL,
                channel1=channel1,
                channel2=channel2,
                thumbnail_channel1=thumbnail_channel1,
                thumbnail_channel2=thumbnail_channel2,
                min_num_blocks=25,
            )
            mr_aligner._coarse_affine_matrix = aligner.coarse_affine_matrix
            mr_aligner.align()
            mr_aligner.constrain_shifts()

            fig = mr_aligner.plot_shifts()
            fig.suptitle(f"{p2.name} (multi-res aligment)", fontsize=8)
            fig.axes[0].set_title(p1.name, fontsize=6)
            save_all_figs(out_dir=out_dir / "qc", format="png", dpi=144)

            pickle_dir = out_dir / "pickle"
            if not pickle_dir.exists():
                pickle_dir.mkdir(parents=True)
            import pickle

            with open(pickle_dir / f"{p2.name}-palom.pkl", "wb") as f:
                pickle.dump(mr_aligner, f)

            aligner = mr_aligner.base_aligner
            block_mx = mr_aligner.base_aligner.block_affine_matrices_da
        elif multi_obj:
            mo_aligner = palom.align_multi_obj.MultiObjAligner(
                r1,
                r2,
                level1=LEVEL,
                channel1=channel1,
                channel2=channel2,
                thumbnail_channel1=thumbnail_channel1,
                thumbnail_channel2=thumbnail_channel2,
            )
            mo_aligner._affine_matrix = aligner.affine_matrix
            mo_aligner._coarse_affine_matrix = aligner.coarse_affine_matrix
            if multi_obj_kwarg is None:
                multi_obj_kwarg = {}
            mo_aligner.run(**multi_obj_kwarg)
            save_all_figs(
                out_dir=out_dir / "qc" / p2.stem, format="png", dpi=144, prefix=p2.name
            )
            block_mx = mo_aligner.block_affine_matrices_da

    if not only_qc:
        mx = aligner.affine_matrix
        if not only_coarse:
            mx = block_mx

        mosaic = palom.align.block_affine_transformed_moving_img(
            ref_img=aligner.ref_img, moving_img=r2.pyramid[LEVEL], mxs=mx
        )
        palom.pyramid.write_pyramid(
            mosaics=[mosaic],
            output_path=out_path,
            pixel_size=r1.pixel_size * r1.level_downsamples[LEVEL],
            channel_names=[list("RBG")],
            compression="zlib",
            downscale_factor=2,
            save_RAM=True,
            tile_size=1024,
            kwargs_tifffile=dict(photometric="rgb", planarconfig="separate"),
        )
    return 0


def viz_coarse(r1, r2, level1, level2, channel1, channel2, mx):
    try:
        import napari
    except ImportError:
        return
    import dask.array as da

    v = napari.Viewer()
    is_bf1 = palom.img_util.is_brightfield_img(r1.pyramid[-1][channel1])
    is_bf2 = palom.img_util.is_brightfield_img(r2.pyramid[-1][channel2])
    inv = {True: da.invert, False: da.array}
    v.add_image(
        [inv[is_bf1](p[channel1]) for p in r1.pyramid[level1:]],
        colormap="bop orange",
        blending="additive",
    )
    v.add_image(
        [inv[is_bf2](p[channel2]) for p in r2.pyramid[level2:]],
        affine=palom.img_util.to_napari_affine(mx),
        colormap="bop blue",
        blending="additive",
    )
    return v


def get_reader(path):
    path = pathlib.Path(path)
    if path.suffix in [".svs", ".ndpi"]:
        return palom.reader.SvsReader
    else:
        return palom.reader.OmePyramidReader


def set_matplotlib_font(font_size=12):
    font_families = matplotlib.rcParams["font.sans-serif"]
    if font_families[0] != "Arial":
        font_families.insert(0, "Arial")
    matplotlib.rcParams["pdf.fonttype"] = 42
    matplotlib.rcParams.update({"font.size": font_size})


def save_all_figs(dpi=300, format="pdf", out_dir=None, prefix=None):
    figs = [plt.figure(i) for i in plt.get_fignums()]
    if prefix is not None:
        for f in figs:
            if f._suptitle:
                f.suptitle(f"{prefix} {f._suptitle.get_text()}")
            else:
                f.suptitle(prefix)
    names = [f._suptitle.get_text() if f._suptitle else "" for f in figs]
    out_dir = pathlib.Path(out_dir)
    out_dir.mkdir(exist_ok=True, parents=True)

    for f, n, nm in zip(figs, plt.get_fignums(), names):
        f.savefig(out_dir / f"{n}-{nm}.{format}", dpi=dpi, bbox_inches="tight")
        plt.close(f)


def set_subplot_size(w, h, ax=None):
    """w, h: width, height in inches"""
    if not ax:
        ax = plt.gca()
    l = ax.figure.subplotpars.left
    r = ax.figure.subplotpars.right
    t = ax.figure.subplotpars.top
    b = ax.figure.subplotpars.bottom
    figw = float(w) / (r - l)
    figh = float(h) / (t - b)
    ax.figure.set_size_inches(figw, figh)


if __name__ == "__main__":
    import fire
    import sys

    fire.Fire(align_he)

    if ("--viz_coarse_napari" in sys.argv) or ("-v" in sys.argv):
        try:
            import napari
        except ImportError:
            print("napari is not installed")
        else:
            napari.run()

    """
    Example 1: inspect coarse alignment using napari
    python align_he.py \
        Z:\RareCyte-S3\P54_CRCstudy_Bridge\P54_S33_Full_Or6_A31_C90c_HMS@20221025_001610_632297.ome.tiff \
        "X:\crc-scans\histowiz scans\20230105-orion_2_cycles\22199$P54_33_HE$US$SCAN$OR$001 _104050.svs" \
        "X:\crc-scans\histowiz scans\20230105-orion_2_cycles\test" \
        --px_size1 0.325 --only_qc --only_coarse --viz_coarse_napari

    Example 2: process pair and output registered image
    python align_he.py \
        Z:\RareCyte-S3\P54_CRCstudy_Bridge\P54_S33_Full_Or6_A31_C90c_HMS@20221025_001610_632297.ome.tiff \
        "X:\crc-scans\histowiz scans\20230105-orion_2_cycles\22199$P54_33_HE$US$SCAN$OR$001 _104050.svs" \
        "X:\crc-scans\histowiz scans\20230105-orion_2_cycles\test" \
        --px_size1 0.325
   
    """
