#!/usr/bin/env python3
"""
Created on Mon Jan  3 19:36:44 2022

@author: varunkapoor
"""

from stardist.models import StarDist2D

from .pretrained import get_registered_models, get_model_details, get_model_instance
import sys


class StarDist2D(StarDist2D):
    def __init__(self, config, name=None, basedir="."):
        super().__init__(config=config, name=name, basedir=basedir)

    @classmethod
    def local_from_pretrained(cls, name_or_alias=None):
        try:
            print(cls)
            get_model_details(cls, name_or_alias, verbose=True)
            return get_model_instance(cls, name_or_alias)
        except ValueError:
            if name_or_alias is not None:
                print(
                    "Could not find model with name or alias '%s'" % (name_or_alias),
                    file=sys.stderr,
                )
                sys.stderr.flush()
            get_registered_models(cls, verbose=True)

    def predict_vollseg(
        self,
        img,
        axes=None,
        normalizer=None,
        prob_thresh=None,
        nms_thresh=None,
        n_tiles=None,
        show_tile_progress=True,
        verbose=False,
        return_labels=True,
        predict_kwargs=None,
        nms_kwargs=None,
        overlap_label=None,
        return_predict=False,
    ):
        """Predict instance segmentation from input image.
        Parameters
        ----------
        img : :class:`numpy.ndarray`
            Input image
        axes : str or None
            Axes of the input ``img``.
            ``None`` denotes that axes of img are the same as denoted in the config.
        normalizer : :class:`csbdeep.data.Normalizer` or None
            (Optional) normalization of input image before prediction.
            Note that the default (``None``) assumes ``img`` to be already normalized.
        sparse: bool
            If true, aggregate probabilities/distances sparsely during tiled
            prediction to save memory (recommended).
        prob_thresh : float or None
            Consider only object candidates from pixels with predicted object probability
            above this threshold (also see `optimize_thresholds`).
        nms_thresh : float or None
            Perform non-maximum suppression that considers two objects to be the same
            when their area/surface overlap exceeds this threshold (also see `optimize_thresholds`).
        n_tiles : iterable or None
            Out of memory (OOM) errors can occur if the input image is too large.
            To avoid this problem, the input image is broken up into (overlapping) tiles
            that are processed independently and re-assembled.
            This parameter denotes a tuple of the number of tiles for every image axis (see ``axes``).
            ``None`` denotes that no tiling should be used.
        show_tile_progress: bool
            Whether to show progress during tiled prediction.
        verbose: bool
            Whether to print some info messages.
        return_labels: bool
            Whether to create a label image, otherwise return None in its place.
        predict_kwargs: dict
            Keyword arguments for ``predict`` function of Keras model.
        nms_kwargs: dict
            Keyword arguments for non-maximum suppression.
        overlap_label: scalar or None
            if not None, label the regions where polygons overlap with that value
        return_predict: bool
            Also return the outputs of :func:`predict` (in a separate tuple)
            If True, implies sparse = False
        Returns
        -------
        (:class:`numpy.ndarray`, dict), (optional: return tuple of :func:`predict`)
            Returns a tuple of the label instances image and also
            a dictionary with the details (coordinates, etc.) of all remaining polygons/polyhedra.
        """
        if predict_kwargs is None:
            predict_kwargs = {}
        if nms_kwargs is None:
            nms_kwargs = {}

        nms_kwargs.setdefault("verbose", verbose)
        _axes = self._normalize_axes(img, axes)
        _axes_net = self.config.axes
        _permute_axes = self._make_permute_axes(_axes, _axes_net)
        _shape_inst = tuple(
            s for s, a in zip(_permute_axes(img).shape, _axes_net) if a != "C"
        )

        res = self.predict(
            img,
            axes=axes,
            normalizer=normalizer,
            n_tiles=n_tiles,
            show_tile_progress=show_tile_progress,
            **predict_kwargs
        )

        res = tuple(res) + (None,)
        prob, dist, points = res
        prob_class = None

        res_instances, polys = self._instances_from_prediction(
            _shape_inst,
            prob,
            dist,
            points=points,
            prob_class=prob_class,
            prob_thresh=prob_thresh,
            nms_thresh=nms_thresh,
            return_labels=return_labels,
            overlap_label=overlap_label,
            **nms_kwargs
        )

        return res_instances, prob, dist
