from __future__ import division

import math

import numpy as np

import cv2

# import and use one of 3 libraries PIL, cv2, or scipy in that order

USE_PIL = True

USE_CV2 = False

USE_SCIPY = False

try:

    import PIL

    from PIL import Image

    raise ImportError

except ImportError:

    USE_PIL = False

if not USE_PIL:

    USE_CV2 = True

    try:

        import cv2

    except ImportError:

        USE_CV2 = False

if not USE_PIL and not USE_CV2:

    USE_SCIPY = True

    try:

        import scipy

        from scipy import misc

    except ImportError:

        USE_SCIPY = False

        raise RuntimeError("couldn't load ANY image library")

class _OtsuPyramid(object):

    """segments histogram into pyramid of histograms, each histogram

    half the size of the previous. Also generate omega and mu values

    for each histogram in the pyramid.

    """

    def load_image(self, im, bins=256):

        """ bins is number of intensity levels """

        if not type(im) == np.ndarray:

            raise ValueError(

                'must be passed numpy array. Got ' + str(type(im)) +

                ' instead'

            )

        if im.ndim == 3:

            raise ValueError(

                'image must be greyscale (and single value per pixel)'

            )

        self.im = im

        hist, ranges = np.histogram(im, bins) #将输入转化为直方图

        # print("hist:",hist,"range:",ranges) # hist表示每个区间内的元素个数，ranges表示区间 

        # convert the numpy array to list of ints

        hist = [int(h) for h in hist]

        histPyr, omegaPyr, muPyr, ratioPyr = \

            self._create_histogram_and_stats_pyramids(hist)

        # arrange so that pyramid[0] is the smallest pyramid

        self.omegaPyramid = [omegas for omegas in reversed(omegaPyr)] # 前缀和

        #  print("self.omeaPtramid:",len(self.omegaPyramid[0]))

        self.muPyramid = [mus for mus in reversed(muPyr)]# 乘了像素的 前缀和

        self.ratioPyramid = ratioPyr# [1, 2, 2, 2, 2, 2, 2, 2]

    def _create_histogram_and_stats_pyramids(self, hist):

        """Expects hist to be a single list of numbers (no numpy array)

        takes an input histogram (with 256 bins) and iteratively

        compresses it by a factor of 2 until the last compressed

        histogram is of size 2. It stores all these generated histograms

        in a list-like pyramid structure. Finally, create corresponding

        omega and mu lists for each histogram and return the 3

        generated pyramids.

        """

        bins = len(hist)

        # eventually you can replace this with a list if you cannot evenly

        # compress a histogram

        ratio = 2

        reductions = int(math.log(bins, ratio)) #ln(bins)/ln(ratio)，约等于开方

        compressionFactor = []

        histPyramid = []

        omegaPyramid = []

        muPyramid = []

        for _ in range(reductions):

            histPyramid.append(hist)

            reducedHist = [sum(hist[i:i+ratio]) for i in range(0, bins, ratio)]

            # collapse a list to half its size, combining the two collpased

            # numbers into one

            hist = reducedHist

            # update bins to reflect the img_nums of the new histogram

            bins = bins // ratio

            compressionFactor.append(ratio)

        # first "compression" was 1, aka it's the original histogram

        compressionFactor[0] = 1

        # print("the length of histPyramid:",len(histPyramid))

        for hist in histPyramid:

            omegas, mus, muT = \

                self._calculate_omegas_and_mus_from_histogram(hist)

            omegaPyramid.append(omegas)

            muPyramid.append(mus)

        return histPyramid, omegaPyramid, muPyramid, compressionFactor

    def _calculate_omegas_and_mus_from_histogram(self, hist):

        """ Comput histogram statistical data: omega and mu for each

        intensity level in the histogram

        """

        probabilityLevels, meanLevels = \

            self._calculate_histogram_pixel_stats(hist)

        bins = len(probabilityLevels)

        # these numbers are critical towards calculations, so we make sure

        # they are float

        ptotal = float(0)

        # sum of probability levels up to k

        omegas = []

        for i in range(bins):

            ptotal += probabilityLevels[i]

            omegas.append(ptotal)

        mtotal = float(0)

        mus = []

        for i in range(bins):

            mtotal += meanLevels[i]

            mus.append(mtotal)

        # muT is the total mean levels.

        muT = float(mtotal)

        return omegas, mus, muT

    def _calculate_histogram_pixel_stats(self, hist):

        """Given a histogram, compute pixel probability and mean

        levels for each bin in the histogram. Pixel probability

        represents the likely-hood that a pixel's intensty resides in

        a specific bin. Pixel mean is the intensity-weighted pixel

        probability.

        """

        # bins = number of intensity levels

        bins = len(hist)

        # print("bins:",bins)

        # N = number of pixels in image. Make it float so that division by

        # N will be a float

        N = float(sum(hist))

        # percentage of pixels at each intensity level: i => P_i

        hist_probability = [hist[i] / N for i in range(bins)]

        # mean level of pixels at intensity level i   => i * P_i

        pixel_mean = [i * hist_probability[i] for i in range(bins)]

        # print("N:",N)

        return hist_probability, pixel_mean

class OtsuFastMultithreshold(_OtsuPyramid):

    """Sacrifices precision for speed. OtsuFastMultithreshold can dial

    in to the threshold but still has the possibility that its

    thresholds will not be the same as a naive-Otsu's method would give

    """

    def calculate_k_thresholds(self, k):

        self.threshPyramid = []

        start = self._get_smallest_fitting_pyramid(k)

        self.bins = len(self.omegaPyramid[start])

        # print("self.bins:",self.bins)

        thresholds = self._get_first_guess_thresholds(k)

        # print("thresholds:",thresholds)

        # give hunting algorithm full range so that initial thresholds

        # can become any value (0-bins)

        deviate = self.bins // 2

        for i in range(start, len(self.omegaPyramid)):

            omegas = self.omegaPyramid[i] # 个数平均前缀和

            mus = self.muPyramid[i] # 平均像素前缀和

            hunter = _ThresholdHunter(omegas, mus, deviate)

            thresholds = \

                hunter.find_best_thresholds_around_estimates(thresholds)

            self.threshPyramid.append(thresholds)

            # how much our "just analyzed" pyramid was compressed from the

            # previous one

            scaling = self.ratioPyramid[i]

            # deviate should be equal to the compression factor of the

            # previous histogram.

            deviate = scaling

            thresholds = [t * scaling for t in thresholds]

        # return readjusted threshold (since it was scaled up incorrectly in

        # last loop)

        # print("true thresholds:",thresholds)

        return [t // scaling for t in thresholds]

    def _get_smallest_fitting_pyramid(self, k):

        """Return the index for the smallest pyramid set that can fit

        K thresholds

        """

        for i, pyramid in enumerate(self.omegaPyramid):

            if len(pyramid) >= k:

                return i

    def _get_first_guess_thresholds(self, k):

        """Construct first-guess thresholds based on number of

        thresholds (k) and constraining intensity values. FirstGuesses

        will be centered around middle intensity value.

        """

        kHalf = k // 2

        midway = self.bins // 2

        firstGuesses = [midway - i for i in range(kHalf, 0, -1)] + [midway] + \

            [midway + i for i in range(1, kHalf)]

        # print("firstGuesses:",firstGuesses)

        # additional threshold in case k is odd

        firstGuesses.append(self.bins - 1)

        return firstGuesses[:k]

    def apply_thresholds_to_image(self, thresholds, im=None):

        if im is None:

            im = self.im

        k = len(thresholds)

        bookendedThresholds = [None] + thresholds + [None]

        # I think you need to use 255 / k *...

        greyValues = [0] + [int(256 / k * (i + 1)) for i in range(0, k - 1)] \

            + [255]

        greyValues = np.array(greyValues, dtype=np.uint8)

        finalImage = np.zeros(im.shape, dtype=np.uint8)

        for i in range(k + 1):

            kSmall = bookendedThresholds[i]

            # True portions of bw represents pixels between the two thresholds

            bw = np.ones(im.shape, dtype=np.bool8)

            if kSmall:

                bw = (im >= kSmall)

            kLarge = bookendedThresholds[i + 1]

            if kLarge:

                bw &= (im < kLarge)

            greyLevel = greyValues[i]

            # apply grey-color to black-and-white image

            greyImage = bw * greyLevel

            # add grey portion to image. There should be no overlap between

            # each greyImage added

            finalImage += greyImage

        return finalImage

class _ThresholdHunter(object):

    """Hunt/deviate around given thresholds in a small region to look

    for a better threshold

    """

    def __init__(self, omegas, mus, deviate=2):

        self.sigmaB = _BetweenClassVariance(omegas, mus)

        # used to be called L

        self.bins = self.sigmaB.bins

        # hunt 2 (or other amount) to either side of thresholds

        self.deviate = deviate

    def find_best_thresholds_around_estimates(self, estimatedThresholds):

        """Given guesses for best threshold, explore to either side of

        the threshold and return the best result.

        """

        bestResults = (

            0, estimatedThresholds, [0 for t in estimatedThresholds]

        )

        # print("bestResults:",bestResults)

        bestThresholds = estimatedThresholds

        bestVariance = 0

        for thresholds in self._jitter_thresholds_generator(

                estimatedThresholds, 0, self.bins):

            # print("thresholds:",thresholds)

            variance = self.sigmaB.get_total_variance(thresholds)

            if variance == bestVariance:

                if sum(thresholds) < sum(bestThresholds):

                    # keep lowest average set of thresholds

                    bestThresholds = thresholds

            elif variance > bestVariance:

                bestVariance = variance

                bestThresholds = thresholds

        return bestThresholds

    def find_best_thresholds_around_estimates_experimental(self, estimatedThresholds):

        """Experimental threshold hunting uses scipy optimize method.

        Finds ok thresholds but doesn't work quite as well

        """

        estimatedThresholds = [int(k) for k in estimatedThresholds]

        if sum(estimatedThresholds) < 10:

            return self.find_best_thresholds_around_estimates(

                estimatedThresholds

            )

        # print('estimated', estimatedThresholds)

        fxn_to_minimize = lambda x: -1 * self.sigmaB.get_total_variance(

            [int(k) for k in x]

        )

        bestThresholds = scipy.optimize.fmin(

            fxn_to_minimize, estimatedThresholds

        )

        bestThresholds = [int(k) for k in bestThresholds]

        # print('bestTresholds', bestThresholds)

        return bestThresholds

    def _jitter_thresholds_generator(self, thresholds, min_, max_):

        pastThresh = thresholds[0]

        # print("pastThresd:",pastThresh)

        if len(thresholds) == 1:

            # -2 through +2

            for offset in range(-self.deviate, self.deviate + 1):

                thresh = pastThresh + offset

                if thresh < min_ or thresh >= max_:

                    # skip since we are conflicting with bounds

                    continue

                yield [thresh]

        else:

            # new threshold without our threshold included

            thresholds = thresholds[1:]

            # number of threshold left to generate in chain

            m = len(thresholds)

            for offset in range(-self.deviate, self.deviate + 1):

                thresh = pastThresh + offset

                # verify we don't use the same value as the previous threshold

                # and also verify our current threshold will not push the last

                # threshold past max

                if thresh < min_ or thresh + m >= max_:

                    continue

                recursiveGenerator = self._jitter_thresholds_generator(

                    thresholds, thresh + 1, max_

                )

                for otherThresholds in recursiveGenerator:

                    yield [thresh] + otherThresholds

class _BetweenClassVariance(object):

    def __init__(self, omegas, mus):

        self.omegas = omegas

        self.mus = mus

        # number of bins / luminosity choices

        self.bins = len(mus)

        self.muTotal = sum(mus)

    def get_total_variance(self, thresholds):

        """Function will pad the thresholds argument with minimum and

        maximum thresholds to calculate between class variance

        """

        thresholds = [0] + thresholds + [self.bins - 1]

        numClasses = len(thresholds) - 1

        # print("thesholds:",thresholds)

        sigma = 0

        for i in range(numClasses):

            k1 = thresholds[i]

            k2 = thresholds[i+1]

            sigma += self._between_thresholds_variance(k1, k2)

        # print("sigma:",sigma)

        return sigma

    def _between_thresholds_variance(self, k1, k2):

        """to be usedin calculating between-class variances only!"""

        # print("len(self.omegas)):",len(self.omegas))

        # print("k2:",k2, "kl:",k1)

        omega = self.omegas[k2] - self.omegas[k1]

        mu = self.mus[k2] - self.mus[k1]

        muT = self.muTotal

        return omega * ((mu - muT)**2)

import math 

def normalize(img):

    # 需要特别注意的是无穷大和无穷小 

    Max = -1

    Min = 10000

    for i in range(img.shape[0]):

        for j in range(img.shape[1]):

            if img[i][j] > Max :

                Max = img[i][j]

            if img[i][j] < Min:

                Min = img[i][j]

    print(Max, Min)

    for i in range(img.shape[0]):

        for j in range(img.shape[0]):

            if math.isnan(img[i][j]):

                img[i][j] = 255.

    img = (img - Min) / (Max - Min)

    return img * 255.

def _otsu(img, categories_pixel_nums = 1):

    '''

    ret_num: the number of image return

    img_nums: the number of calculate thresholds

    categories_pixel_nums: the number of pixel categories

    return the images with only categories_pixel_nums types of pixels

    '''

    return normalize(img.copy())

    # img = normalize(img.copy())

    # ot = OtsuFastMultithreshold()

    # ot.load_image(img)

    # kThresholds = ot.calculate_k_thresholds(categories_pixel_nums)

    # #  print(total)

    # return ot.apply_thresholds_to_image(kThresholds,img)

    # plt.imsave('C:/Users/RL/Desktop/可解释性的特征学习分类/otsu/' + str(i) + '.jpg',crushed,cmap="gray")

    # ans.append(crushed)

    # plt.figure()

    # plt.subplot(121)

    # plt.imshow(img[i],cmap="gray")

    # plt.subplot(122)

    # plt.imshow(crushed,cmap="gray")

    # plt.show()

#otsu的接口，返回处理后的图片

def otsu_helper(img, upper=0.5, down = -0.5,categories=1):

    ot = OtsuFastMultithreshold()

    ot.load_image(img)

    return ot.apply_thresholds_to_image([down, upper], img)