from pathlib import Path
from pkg_resources import get_distribution
from zipfile import ZipFile
import numpy as np
import tempfile
from distutils.version import LooseVersion
from csbdeep.utils import axes_check_and_normalize, normalize, _raise


DEEPIMAGEJ_MACRO = \
"""
//*******************************************************************
// Date: July-2021
// Credits: StarDist, DeepImageJ
// URL:
//      https://github.com/stardist/stardist
//      https://deepimagej.github.io/deepimagej
// This macro was adapted from
// https://github.com/deepimagej/imagej-macros/blob/648caa867f6ccb459649d4d3799efa1e2e0c5204/StarDist2D_Post-processing.ijm
// Please cite the respective contributions when using this code.
//*******************************************************************
//  Macro to run StarDist postprocessing on 2D images.
//  StarDist and deepImageJ plugins need to be installed.
//  The macro assumes that the image to process is a stack in which
//  the first channel corresponds to the object probability map
//  and the remaining channels are the radial distances from each
//  pixel to the object boundary.
//*******************************************************************

// Get the name of the image to call it
getDimensions(width, height, channels, slices, frames);
name=getTitle();

probThresh={probThresh};
nmsThresh={nmsThresh};

// Isolate the detection probability scores
run("Make Substack...", "channels=1");
rename("scores");

// Isolate the oriented distances
run("Fire");
selectWindow(name);
run("Delete Slice", "delete=channel");
selectWindow(name);
run("Properties...", "channels=" + maxOf(channels, slices) - 1 + " slices=1 frames=1 pixel_width=1.0000 pixel_height=1.0000 voxel_depth=1.0000");
rename("distances");
run("royal");

// Run StarDist plugin
run("Command From Macro", "command=[de.csbdresden.stardist.StarDist2DNMS], args=['prob':'scores', 'dist':'distances', 'probThresh':'" + probThresh + "', 'nmsThresh':'" + nmsThresh + "', 'outputType':'Both', 'excludeBoundary':'2', 'roiPosition':'Stack', 'verbose':'false'], process=[false]");
"""


def _import(error=True):
    try:
        from importlib_metadata import metadata
        from bioimageio.core.build_spec import build_model # type: ignore
        import xarray as xr
        import bioimageio.core # type: ignore
    except ImportError:
        if error:
            raise RuntimeError(
                "Required libraries are missing for bioimage.io model export.\n"
                "Please install StarDist as follows: pip install 'stardist[bioimageio]'\n"
                "(You do not need to uninstall StarDist first.)"
            )
        else:
            return None
    return metadata, build_model, bioimageio.core, xr


def _create_stardist_dependencies(outdir):
    from ruamel.yaml import YAML
    from tensorflow import __version__ as tf_version
    from . import __version__ as stardist_version
    pkg_info = get_distribution("stardist")
    # dependencies that start with the name "bioimageio" will be added as conda dependencies
    reqs_conda = [str(req) for req in pkg_info.requires(extras=['bioimageio']) if str(req).startswith('bioimageio')]
    # only stardist and tensorflow as pip dependencies
    tf_major, tf_minor = LooseVersion(tf_version).version[:2]
    reqs_pip = (f"stardist>={stardist_version}", f"tensorflow>={tf_major}.{tf_minor},<{tf_major+1}")
    # conda environment
    env = dict(
        name = 'stardist',
        channels = ['defaults', 'conda-forge'],
        dependencies = [
            ('python>=3.7,<3.8' if tf_major == 1 else 'python>=3.7'),
            *reqs_conda,
            'pip', {'pip': reqs_pip},
        ],
    )
    yaml = YAML(typ='safe')
    path = outdir / "environment.yaml"
    with open(path, "w") as f:
        yaml.dump(env, f)
    return f"conda:{path}"


def _create_stardist_doc(outdir):
    doc_path = outdir / "README.md"
    text = (
        "# StarDist Model\n"
        "This is a model for object detection with star-convex shapes.\n"
        "Please see the [StarDist repository](https://github.com/stardist/stardist) for details."
    )
    with open(doc_path, "w") as f:
        f.write(text)
    return doc_path


def _get_stardist_metadata(outdir, model):
    metadata, *_ = _import()
    package_data = metadata("stardist")
    doi_2d = "https://doi.org/10.1007/978-3-030-00934-2_30"
    doi_3d = "https://doi.org/10.1109/WACV45572.2020.9093435"
    authors = {
        'Martin Weigert': dict(name='Martin Weigert', github_user='maweigert'),
        'Uwe Schmidt': dict(name='Uwe Schmidt', github_user='uschmidt83'),
    }
    data = dict(
        description=package_data["Summary"],
        authors=list(authors.get(name.strip(),dict(name=name.strip())) for name in package_data["Author"].split(",")),
        git_repo=package_data["Home-Page"],
        license=package_data["License"],
        dependencies=_create_stardist_dependencies(outdir),
        cite=[{"text": "Cell Detection with Star-Convex Polygons", "doi": doi_2d},
              {"text": "Star-convex Polyhedra for 3D Object Detection and Segmentation in Microscopy", "doi": doi_3d}],
        tags=[
            'fluorescence-light-microscopy', 'whole-slide-imaging', 'other', # modality
            f'{model.config.n_dim}d', # dims
            'cells', 'nuclei', # content
            'tensorflow', # framework
            'fiji', # software
            'unet', # network
            'instance-segmentation', 'object-detection', # task
            'stardist',
        ],
        covers=["https://raw.githubusercontent.com/stardist/stardist/master/images/stardist_logo.jpg"],
        documentation=_create_stardist_doc(outdir),
    )
    return data


def _predict_tf(model_path, test_input):
    import tensorflow as tf
    from csbdeep.utils.tf import IS_TF_1
    # need to unzip the model assets
    model_assets = model_path.parent / "tf_model"
    with ZipFile(model_path, "r") as f:
        f.extractall(model_assets)
    if IS_TF_1:
        # make a new graph, i.e. don't use the global default graph
        with tf.Graph().as_default():
            with tf.Session() as sess:
                tf_model = tf.saved_model.load_v2(str(model_assets))
                x = tf.convert_to_tensor(test_input, dtype=tf.float32)
                model = tf_model.signatures["serving_default"]
                y = model(x)
                sess.run(tf.global_variables_initializer())
                output = sess.run(y["output"])
    else:
        tf_model = tf.saved_model.load(str(model_assets))
        x = tf.convert_to_tensor(test_input, dtype=tf.float32)
        model = tf_model.signatures["serving_default"]
        y = model(x)
        output = y["output"].numpy()
    return output


def _get_weights_and_model_metadata(outdir, model, test_input, test_input_axes, test_input_norm_axes, mode, min_percentile, max_percentile):

    # get the path to the exported model assets (saved in outdir)
    if mode == "keras_hdf5":
        raise NotImplementedError("Export to keras format is not supported yet")
    elif mode == "tensorflow_saved_model_bundle":
        assets_uri = outdir / "TF_SavedModel.zip"
        model_csbdeep = model.export_TF(assets_uri, single_output=True, upsample_grid=True)
    else:
        raise ValueError(f"Unsupported mode: {mode}")

    # to force "inputs.data_type: float32" in the spec (bonus: disables normalization warning in model._predict_setup)
    test_input = test_input.astype(np.float32)

    # convert test_input to axes_net semantics and shape, also resize if necessary (to adhere to axes_net_div_by)
    test_input, axes_img, axes_net, axes_net_div_by, *_ = model._predict_setup(
        img=test_input,
        axes=test_input_axes,
        normalizer=None,
        n_tiles=None,
        show_tile_progress=False,
        predict_kwargs={},
    )

    # normalization axes string and numeric indices
    axes_norm = set(axes_net).intersection(set(axes_check_and_normalize(test_input_norm_axes, disallowed='S')))
    axes_norm = "".join(a for a in axes_net if a in axes_norm)  # preserve order of axes_net
    axes_norm_num = tuple(axes_net.index(a) for a in axes_norm)

    # normalize input image
    test_input_norm = normalize(test_input, pmin=min_percentile, pmax=max_percentile, axis=axes_norm_num)

    net_axes_in = axes_net.lower()
    net_axes_out = axes_check_and_normalize(model._axes_out).lower()
    ndim_tensor = len(net_axes_out) + 1

    input_min_shape = list(axes_net_div_by)
    input_min_shape[axes_net.index('C')] = model.config.n_channel_in
    input_step = list(axes_net_div_by)
    input_step[axes_net.index('C')] = 0

    # add the batch axis to shape and step
    input_min_shape = [1] + input_min_shape
    input_step = [0] + input_step

    # the axes strings in bioimageio convention
    input_axes = "b" + net_axes_in.lower()
    output_axes = "b" + net_axes_out.lower()

    if mode == "keras_hdf5":
        output_names = ("prob", "dist") + (("class_prob",) if model._is_multiclass() else ())
        output_n_channels = (1, model.config.n_rays,) + ((1,) if model._is_multiclass() else ())
        # the output shape is computed from the input shape using
        # output_shape[i] = output_scale[i] * input_shape[i] + 2 * output_offset[i]
        output_scale = [1]+list(1/g for g in model.config.grid) + [0]
        output_offset = [0]*(ndim_tensor)

    elif mode == "tensorflow_saved_model_bundle":
        if model._is_multiclass():
            raise NotImplementedError("Tensorflow SavedModel not supported for multiclass models yet")
        # regarding input/output names: https://github.com/CSBDeep/CSBDeep/blob/b0d2f5f344ebe65a9b4c3007f4567fe74268c813/csbdeep/utils/tf.py#L193-L194
        input_names = ["input"]
        output_names = ["output"]
        output_n_channels = (1 + model.config.n_rays,)
        # the output shape is computed from the input shape using
        # output_shape[i] = output_scale[i] * input_shape[i] + 2 * output_offset[i]
        # same shape as input except for the channel dimension
        output_scale = [1]*(ndim_tensor)
        output_scale[output_axes.index("c")] = 0
        # no offset, except for the input axes, where it is output channel / 2
        output_offset = [0.0]*(ndim_tensor)
        output_offset[output_axes.index("c")] = output_n_channels[0] / 2.0

    assert all(s in (0, 1) for s in output_scale), "halo computation assumption violated"
    halo = model._axes_tile_overlap(output_axes.replace('b', 's'))
    halo = [int(np.ceil(v/8)*8) for v in halo]  # optional: round up to be divisible by 8

    # the output shape needs to be valid after cropping the halo, so we add the halo to the input min shape
    input_min_shape = [ms + 2 * ha for ms, ha in zip(input_min_shape, halo)]

    # make sure the input min shape is still divisible by the min axis divisor
    input_min_shape = input_min_shape[:1] + [ms + (-ms % div_by) for ms, div_by in zip(input_min_shape[1:], axes_net_div_by)]
    assert all(ms % div_by == 0 for ms, div_by in zip(input_min_shape[1:], axes_net_div_by))

    metadata, *_ = _import()
    package_data = metadata("stardist")
    is_2D = model.config.n_dim == 2

    weights_file = outdir / "stardist_weights.h5"
    model.keras_model.save_weights(str(weights_file))

    config = dict(
        stardist=dict(
            python_version=package_data["Version"],
            thresholds=dict(model.thresholds._asdict()),
            weights=weights_file.name,
            config=vars(model.config),
        )
    )

    if is_2D:
        macro_file = outdir / "stardist_postprocessing.ijm"
        with open(str(macro_file), 'w', encoding='utf-8') as f:
            f.write(DEEPIMAGEJ_MACRO.format(probThresh=model.thresholds.prob, nmsThresh=model.thresholds.nms))
        config['stardist'].update(postprocessing_macro=macro_file.name)

    n_inputs = len(input_names)
    assert n_inputs == 1
    input_config = dict(
        input_names=input_names,
        input_min_shape=[input_min_shape],
        input_step=[input_step],
        input_axes=[input_axes],
        input_data_range=[["-inf", "inf"]],
        preprocessing=[[dict(
            name="scale_range",
            kwargs=dict(
                mode="per_sample",
                axes=axes_norm.lower(),
                min_percentile=min_percentile,
                max_percentile=max_percentile,
            ))]]
        )

    n_outputs = len(output_names)
    output_config = dict(
        output_names=output_names,
        output_data_range=[["-inf", "inf"]] * n_outputs,
        output_axes=[output_axes] * n_outputs,
        output_reference=[input_names[0]] * n_outputs,
        output_scale=[output_scale] * n_outputs,
        output_offset=[output_offset] * n_outputs,
        halo=[halo] * n_outputs
    )

    in_path = outdir / "test_input.npy"
    np.save(in_path, test_input[np.newaxis])

    if mode == "tensorflow_saved_model_bundle":
        test_outputs = _predict_tf(assets_uri, test_input_norm[np.newaxis])
    else:
        test_outputs = model.predict(test_input_norm)

    # out_paths = []
    # for i, out in enumerate(test_outputs):
    #     p = outdir / f"test_output{i}.npy"
    #     np.save(p, out)
    #     out_paths.append(p)
    assert n_outputs == 1
    out_paths = [outdir / "test_output.npy"]
    np.save(out_paths[0], test_outputs)

    from tensorflow import __version__ as tf_version
    data = dict(weight_uri=assets_uri, test_inputs=[in_path], test_outputs=out_paths,
                config=config, tensorflow_version=tf_version)
    data.update(input_config)
    data.update(output_config)
    _files = [str(weights_file)]
    if is_2D:
        _files.append(str(macro_file))
    data.update(attachments=dict(files=_files))

    return data


def export_bioimageio(
    model,
    outpath,
    test_input,
    test_input_axes=None,
    test_input_norm_axes='ZYX',
    name=None,
    mode="tensorflow_saved_model_bundle",
    min_percentile=1.0,
    max_percentile=99.8,
    overwrite_spec_kwargs=None,
):
    """Export stardist model into bioimage.io format, https://github.com/bioimage-io/spec-bioimage-io.

    Parameters
    ----------
    model: StarDist2D, StarDist3D
        the model to convert
    outpath: str, Path
        where to save the model
    test_input: np.ndarray
        input image for generating test data
    test_input_axes: str or None
        the axes of the test input, for example 'YX' for a 2d image or 'ZYX' for a 3d volume
        using None assumes that axes of test_input are the same as those of model
    test_input_norm_axes: str
        the axes of the test input which will be jointly normalized, for example 'ZYX' for all spatial dimensions ('Z' ignored for 2D input)
        use 'ZYXC' to also jointly normalize channels (e.g. for RGB input images)
    name: str
        the name of this model (default: None)
        if None, uses the (folder) name of the model (i.e. `model.name`)
    mode: str
        the export type for this model (default: "tensorflow_saved_model_bundle")
    min_percentile: float
        min percentile to be used for image normalization (default: 1.0)
    max_percentile: float
        max percentile to be used for image normalization (default: 99.8)
    overwrite_spec_kwargs: dict or None
        spec keywords that should be overloaded (default: None)
    """
    _, build_model, *_ = _import()
    from .models import StarDist2D, StarDist3D
    isinstance(model, (StarDist2D, StarDist3D)) or _raise(ValueError("not a valid model"))
    0 <= min_percentile < max_percentile <= 100 or _raise(ValueError("invalid percentile values"))

    if name is None:
        name = model.name
    name = str(name)

    outpath = Path(outpath)
    if outpath.suffix == "":
        outdir = outpath
        zip_path = outdir / f"{name}.zip"
    elif outpath.suffix == ".zip":
        outdir = outpath.parent
        zip_path = outpath
    else:
        raise ValueError(f"outpath has to be a folder or zip file, got {outpath}")
    outdir.mkdir(exist_ok=True, parents=True)

    with tempfile.TemporaryDirectory() as _tmp_dir:
        tmp_dir = Path(_tmp_dir)
        kwargs = _get_stardist_metadata(tmp_dir, model)
        model_kwargs = _get_weights_and_model_metadata(tmp_dir, model, test_input, test_input_axes, test_input_norm_axes, mode,
                                                       min_percentile=min_percentile, max_percentile=max_percentile)
        kwargs.update(model_kwargs)
        if overwrite_spec_kwargs is not None:
            kwargs.update(overwrite_spec_kwargs)

        build_model(name=name, output_path=zip_path, add_deepimagej_config=(model.config.n_dim==2), root=tmp_dir, **kwargs)
        print(f"\nbioimage.io model with name '{name}' exported to '{zip_path}'")


def import_bioimageio(source, outpath):
    """Import stardist model from bioimage.io format, https://github.com/bioimage-io/spec-bioimage-io.

    Load a model in bioimage.io format from the given `source` (e.g. path to zip file, URL)
    and convert it to a regular stardist model, which will be saved in the folder `outpath`.

    Parameters
    ----------
    source: str, Path
        bioimage.io resource (e.g. path, URL)
    outpath: str, Path
        folder to save the stardist model (must not exist previously)

    Returns
    -------
    StarDist2D or StarDist3D
        stardist model loaded from `outpath`

    """
    import shutil, uuid
    from csbdeep.utils import save_json
    from .models import StarDist2D, StarDist3D
    *_, bioimageio_core, _ = _import()

    outpath = Path(outpath)
    not outpath.exists() or _raise(FileExistsError(f"'{outpath}' already exists"))

    with tempfile.TemporaryDirectory() as _tmp_dir:
        tmp_dir = Path(_tmp_dir)
        # download the full model content to a temporary folder
        zip_path = tmp_dir / f"{str(uuid.uuid4())}.zip"
        bioimageio_core.export_resource_package(source, output_path=zip_path)
        with ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(tmp_dir)
        zip_path.unlink()
        rdf_path = tmp_dir / "rdf.yaml"
        biomodel = bioimageio_core.load_resource_description(rdf_path)

        # read the stardist specific content
        'stardist' in biomodel.config or _raise(RuntimeError("bioimage.io model not compatible"))
        config = biomodel.config['stardist']['config']
        thresholds = biomodel.config['stardist']['thresholds']
        weights = biomodel.config['stardist']['weights']

        # make sure that the keras weights are in the attachments
        weights_file = None
        for f in biomodel.attachments.files:
            if f.name == weights and f.exists():
                weights_file = f
                break
        weights_file is not None or _raise(FileNotFoundError(f"couldn't find weights file '{weights}'"))

        # save the config and threshold to json, and weights to hdf5 to enable loading as stardist model
        # copy bioimageio files to separate sub-folder
        outpath.mkdir(parents=True)
        save_json(config, str(outpath / 'config.json'))
        save_json(thresholds, str(outpath / 'thresholds.json'))
        shutil.copy(str(weights_file), str(outpath / "weights_bioimageio.h5"))
        shutil.copytree(str(tmp_dir), str(outpath / "bioimageio"))

    model_class = (StarDist2D if config['n_dim'] == 2 else StarDist3D)
    model = model_class(None, outpath.name, basedir=str(outpath.parent))

    return model