import cv2
import matplotlib.pyplot as plt
import peakutils
from functools import wraps

from numpy.random import random

from qrdecode import Scanner
from lmfit.models import *
from scipy.optimize import curve_fit

PEAK_FUNCTIONS = ('gaussian', 'lorentzian', 'voigt', 'pvoigt', 'moffat', 'pearson7')


def gaussian(x, a, b, c):
    """
    Gaussian function
    :param x: profile length
    :param a: heigth of the curve's peak
    :param b: position of the peak
    :param c: full with at half maximum
    :return: y values for each x of the gaussian with a,b and c parameters
    """
    return a * np.exp(-(x - b) ** 2 / (2 * c ** 2))


def func_to_opt(x, a0, b0, c0, a1, b1, c1):
    """
    Adjusts to 2 bell curve shape
    :return:
    :param x: x values, profile length
    :param a0: the height of the first curve's peak
    :param b0: position of the first center peak
    :param c0: full width at half maximum of the first bell
    :param a1: the height of the second curve's peak
    :param b1: position of the second center peak
    :param c1: full width at half maximum of the second bell
    :return:
    """
    return lineshapes.gaussian(x, a0, b0, c0) + lineshapes.gaussian(x, a1, b1, c1)


def calc_mu_sigma(x, a0, b0, c0, a1, b1, c1):
    """

    :param x:
    :param a0:
    :param b0:
    :param c0:
    :param a1:
    :param b1:
    :param c1:
    :return:
    """
    # First gaussian prob
    fprob = lineshapes.gaussian(x, a0, b0, c0) / lineshapes.gaussian(x, a0, b0, c0).sum()
    fmu = x.dot(fprob)
    fmom2 = np.power(x, 2).dot(fprob)  # 2nd moment
    fvar = fmom2 - fmu ** 2  # variance
    fstd = np.sqrt(fvar)

    # Second gaussian prob
    sprob = lineshapes.gaussian(x, a1, b1, c1) / lineshapes.gaussian(x, a1, b1, c1).sum()
    smu = x.dot(sprob)
    smom2 = np.power(x, 2).dot(sprob)
    svar = smom2 - smu ** 2
    sstd = np.sqrt(svar)
    print(fmu, smu)
    return [fmu, fstd, smu, sstd]


class Biostrip:

    def __init__(self, img, **kwargs):
        self.scanner = Scanner(img, **kwargs)
        self.dst = None
        self.src_corners = None
        self.max_height = None
        self.max_width = None
        self.sliced = None

        self.scanner.scan()
        self.strip_scan()
        self.strip_slice()

    @property
    def codes_position(self):
        """
        Returns the codes position in the image found by the scanner.

        :return: The codes list.
        :rtype: string

        """
        return self.scanner.codes_position

    @property
    def image(self):
        """
        Returns the image being used by the scanner.

        :return: The image being used by the scanner.
        :rtype: string

        """
        return self.scanner.image

    @property
    def slice(self):
        """
        Returns the sliced image being generated by strip_slice method.

        :return: The sliced image.
        :rtype: string

        """
        return self.sliced

    def add_image(self, file):
        """
        Adds an image to the scanner.

        :param file: The image to be used.
        :type file: string
        :return: Void function.
        :rtype: None

        """
        self.scanner.add_image(file)

    @property
    def output_border(self):
        """
        The output border.

        :setter: Sets the output border.
        :getter: Returns the output border.
        :type: int or (numpy.ndarray, list, tuple).
        :return: The output border.
        :rtype: numpy.ndarray

        """
        return self.scanner.output_border

    @property
    def output_size(self):
        """
        The output size.

        :setter: Sets the output size.
        :getter: Returns the output size.
        :type: int or (numpy.ndarray, list, tuple).
        :return: The output size.
        :rtype: numpy.ndarray

        """
        return self.scanner.output_size

    @property
    def output_offset(self):
        """
        The output offset.

        :setter: Sets the output offset.
        :getter: Returns the output offset.
        :type: int or (numpy.ndarray, list, tuple).
        :return: The output size.
        :rtype: numpy.ndarray

        """
        return self.scanner.output_offset

    def strip_slice(self):
        """
        Method that slices the strip scanned image obtaining the ROI between de Qrs removing padding and margins from
        sides.
        :return: the sliced image
        """
        remove_top = (self.codes_position[0][0][1] - 10) * 0.4
        remove_bottom = (self.codes_position[0][0][1] - 10) * 0.4

        self.sliced = self.image[int(self.codes_position[1][0][1] + remove_top):int(
            self.codes_position[0][0][1] - 10 - remove_bottom),
                      int(self.codes_position[1][1][0] + 90):int(self.codes_position[0][3][0] - 90)]

        img = cv2.rectangle(self.image.copy(),
                            (int(self.codes_position[1][1][0] + 30), int(self.codes_position[1][0][1] + remove_top)), (
                                int(self.codes_position[0][3][0] - 30),
                                int(self.codes_position[0][0][1] - 10 - remove_bottom)), (0, 0, 255), 3)

        return self.sliced

    def strip_scan(self):
        """
        Prerequisites: As the object scanning is a strip at least 2qrs has been detected with scan method from
        scanner. Takes de qr detected and slices the ROI doing the perspective with the warp matrix created with the
        top-right, top-left, bottom-right and bottom-left points defined by both qrs.
        """
        # now that we have our contour, we need to determine

        # the top-left, top-right, bottom-right, and bottom-left
        # points so that we can later warp the image -- we'll start
        # by reshaping our contour to be our finals and initializing
        # our output rectangle in top-left, top-right, bottom-right,
        # and bottom-left order

        rect = np.zeros((4, 2), dtype="float32")

        # the top-left point has the smallest sum whereas the
        # bottom-right has the largest sum

        rect[0] = self.codes_position[0][2]  # top-left
        rect[2] = self.codes_position[1][2]  # bottom-right

        rect[1] = self.codes_position[1][0]
        rect[3] = self.codes_position[0][0]
        self.src_corners = rect

        # now that we have our rectangle of points, let's compute
        # the width of our new image
        (tl, tr, br, bl) = rect
        width_a = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
        width_b = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))

        # ...and now for the height of our new image
        height_a = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
        height_b = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))

        # take the maximum of the width and height values to reach
        # our final dimensions

        self.max_width = max(int(width_a), int(width_b))
        self.max_height = max(int(height_a), int(height_b))

        # construct our destination points which will be used to
        # map the screen to a top-down, "birds eye" view
        self.dst = np.array([[self.output_border[0], self.max_height],
                             [self.max_width, self.max_height],
                             [self.max_width, self.output_border[1]],
                             self.output_border], dtype="float32")

        # OpenCV solution for performing the transform.
        image_trimmed = cv2.warpPerspective(
            self.image,
            cv2.getPerspectiveTransform(self.src_corners, self.dst),
            (self.max_width + self.output_border[0], self.max_height + self.scanner.output_border[1]),
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(255, 255, 255)
        )

        self.add_image(image_trimmed)
        self.scanner.get_contours()
        self.scanner.scan()  # re scan them
        return image_trimmed

    def strip_back(self):
        """

        :return:
        """
        roi = cv2.imread('/home/alicia/ProjectQR/qrdecode/scripts/images/lila2.png')
        hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
        target = self.image
        plt.figure()
        plt.imshow(target)
        plt.figure()
        hsvt = cv2.cvtColor(target, cv2.COLOR_BGR2HSV)
        # calculating object histogram
        roihist = cv2.calcHist([hsv], [0, 1], None, [180, 256], [0, 180, 0, 256])
        # normalize histogram and apply back-projection
        cv2.normalize(roihist, roihist, 0, 255, cv2.NORM_MINMAX)
        cv2.imshow('hist', roihist)
        plt.imshow(hsvt)
        dst = cv2.calcBackProject([hsvt], [0, 1], roihist, [0, 180, 0, 256], 1)
        # Now convolve with circular disc
        disc = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        cv2.filter2D(dst, -1, disc, dst)
        # threshold and binary AND
        ret, thresh = cv2.threshold(dst, 50, 255, 0)
        thresh = cv2.merge((thresh, thresh, thresh))
        res = cv2.bitwise_and(target, thresh)
        res = np.vstack((target, thresh, res))
        cv2.imshow('result', res)
        return res


class ColorLines:

    PEAK_FUNCTIONS = {}

    def __init__(self, img, model="pvoigt"):
        """
        Sets the image and initializes the necessary saved parameters
        :param img: the ROI subtracted from the image containing the corresponding color lines columns
        :param kwargs:
        """

        self.image = img  # ROI image
        self.meanMat = None
        self._fitting_functions = {}
        self.profile_gray = None
        self.profile = None
        self.profile_old = None
        self.baseline_values = None
        self.fitting_config = None
        self.result = None
        self.model_sel = model

        self.add_fitting_function("gaussian", GaussianModel)
        self.add_fitting_function("lorentzian", LorentzianModel)
        self.add_fitting_function("pearson7", Pearson7Model)
        self.add_fitting_function("pvoigt", PseudoVoigtModel)
        self.add_fitting_function("voigt", VoigtModel)
        self.add_fitting_function("students_t", StudentsTModel)
        self.add_fitting_function("skewed_gaussian", SkewedGaussianModel)
        self.add_fitting_function("moffat", MoffatModel)

    def add_fitting_function(self, name, peak_fitting_func):
        """
        Adds a fitting function to the list of available ones.

        :param peak_fitting_func:
        :param name: The name of the fitting function
        :type name: string
        :return: Void function.
        :rtype: None

        """
        self.PEAK_FUNCTIONS[name] = peak_fitting_func

    def peak_like_model(self, *args, **kwargs):
        """
        Returns the function associated to the fitting function "name"`.

        :param name: The name of the fitting function.
        :type name: string
        :param args: Arguments of the fitting function that will be passed to \
            ``peak_function``.
        :type args: tuple
        :param kwargs: Arguments of the function associated to the threshold that will be passed to \
            ``peak_function``.
        :type kwargs: dict
        :return: The function of fitting (if it exists) or None (if it does not).
        :rtype: function or None

        """
        if self.model_sel in self.PEAK_FUNCTIONS:
            peak_function = self.PEAK_FUNCTIONS[self.model_sel]
            #profile_func = peak_function(self.profile_gray, *args, **kwargs)
            #self._fitting_functions[self.model_sel] = profile_func
            return peak_function(*args, **kwargs)
        else:
            return None

    def mean(self):
        """
        Method that reduces the image to 1xMx3 matrix that the represents the mean of the columns and then passes to
        gray function.
        :return: an 1xM matrix where each column represents the mean of the three channels of the
        passed matrix
        """
        # Mean by columns
        self.meanMat = cv2.reduce(self.image, 0, cv2.REDUCE_AVG, dtype=cv2.CV_32S)

        # To gray
        self._to_gray(self.meanMat[0])

        return self.meanMat

    def containing_colors(self):
        lines = self.contour_line_marks()
        a, b, c = lines.shape
        print(lines.shape)
        for i in range(a):
            cv2.line(self.image, (lines[i][0][0], lines[i][0][1]), (lines[i][0][2], lines[i][0][3]), (0, 0, 255), 3)
        cv2.imshow('Lines Detected', self.image)

    def contour_line_marks(self):
        """
        Method that performs a HOG for searching strip lines.
        :return: the sliced image with the found lines
        """
        gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)

        lines = cv2.HoughLinesP(image=edges, rho=1, theta=np.pi, threshold=10, lines=np.array([]),
                                minLineLength=10, maxLineGap=0.1)

        a, b, c = lines.shape
        for i in range(a):
            cv2.line(self.image, (lines[i][0][0], lines[i][0][1]), (lines[i][0][2], lines[i][0][3]), (0, 0, 255), 3)

        # TODO: Separate two lines points with k-means
        cv2.imshow(self.image)

        return self.image

    def hsv_histogram(self):
        """
        Method that shows the amount of each color component from the sliced image.
        """
        color = ('b', 'g', 'r')
        for i, col in enumerate(color):
            histr = cv2.calcHist([self.image], [i], None, [256], [0, 256])
            plt.plot(histr, color=col)
            plt.xlim([0, 256])
        plt.show()

    def _to_gray(self, im):
        """
        Returns a Grayscale image from an RGB image.
        :param im: An RGB image with shape (N, M, 3).
        :return: A Grayscale image with shape (N, M).
        """
        # TODO different type of weights.
        gray_weights = [1 / 3] * 3
        self.profile_gray = np.dot(self.meanMat, gray_weights)

        # Flipping graph
        self.profile = 100 * (1 - self.profile_gray / 255)

        # Taking the only column as a numpy
        self.profile = self.profile[0]

        return self.profile

    def to_hsv(self):
        # Mean by columns


        hsv_im = cv2.cvtColor(self.image, cv2.COLOR_RGB2HSV)
        self.meanMat = cv2.reduce(hsv_im, 0, cv2.REDUCE_AVG, dtype=cv2.CV_32S)
        #self.profile = cv2.normalize(self.meanMat[:,:,1][0], dst=None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC3)
        self.profile = self.meanMat[:,:,1][0]


    def integrate(self, fmu, fstd, smu, sstd):
        """
        Method that shows the area bellow the curve. Assumes the profile has been extracted.
        Uses trapezoidal method from numpy.
        """
        # Compute the area using the composite trapezoidal rule.
        farea = np.trapz(self.profile[int(fmu - 3 * fstd):int(fmu + 3 * fstd)], dx=5)
        sarea = np.trapz(self.profile[int(smu - 3 * sstd):int(smu + 3 * sstd)], dx=5)
        print("farea = ", farea)
        print("sarea = ", sarea)
        print("h2/h1 = ", sarea/farea)

    def solve(self, fitting ="gaussian"):
        """
        Method that finds baseline, subtracts it from the profile, finds the peaks and fits de gaussian function.
        Then plots the results.
        """
        self.profile_old = self.profile.copy()
        x = np.arange(len(self.profile))
        self.baseline_values = peakutils.baseline(self.profile_old, deg=4, tol=1e-2, max_it=10)

        # calculate baseline y values
        self.profile = self.profile_old - self.baseline_values

        threshold = 0.25 * self.profile.max()

        # Finding peaks with peaksutils
        indexes = peakutils.indexes(self.profile, thres=0.5, min_dist=5)

        peaks_loc = threshold < self.profile
        peaks_ii = np.where(peaks_loc)[0]

        # Split peaks checking discontinuities in indices:
        peaks_collection = [[peaks_ii[0] - 1]]
        peak_i = 0

        for index in peaks_ii:
            # TODO: that's sad method :(
            if index == 1 + peaks_collection[peak_i][-1]:
                peaks_collection[peak_i].append(index)
            else:
                peak_i += 1
                peaks_collection.append([index])

        # Let's compute the reference values!
        x = np.arange(len(self.profile))

        # Guessing the Gaussian peaks:
        a, b, c = [], [], []

        # TODO take the highest peaks of peaks_collection, not the two first..., and use peakutils to improve the values
        for peak in peaks_collection[0:2]:
            # c (gaussian width) is given by the difference in indexes:
            c.append(0.5 * (peak[-1] - peak[0]))

            # b (x offset) is given by the mean index:
            b.append(int(np.mean(peak)))

            # a (amplitude) is given by the maximum value substracting the baseline
            peak_values = self.profile[peak]
            a.append(peak_values.max())

        lmodel1 = self.peak_like_model(prefix="f")
        lmodel2 = self.peak_like_model(prefix="ss")

        lmodel1.set_param_hint('famplitude', value=a[0])
        lmodel1.set_param_hint('fcenter', value=b[0])
        lmodel1.set_param_hint('fsigma', value=c[0])

        par1 = lmodel1.make_params(famplitude=a[0], fcenter=b[0], fsigma=c[0])

        lmodel2.set_param_hint('ssamplitude', value=a[1])
        lmodel2.set_param_hint('sscenter', value=b[1])
        lmodel2.set_param_hint('sssigma', value=c[1])

        par2 = lmodel2.make_params(ssamplitude=a[1], sscenter=b[1], sssigma=c[1])

        par = par2 + par1
        lmodel = lmodel1 + lmodel2
        self.result = lmodel.fit(self.profile, parms=par, x=x)
        print(self.result.values['ssheight']/self.result.values['ssheight'])

        # popt, pcov = curve_fit(func_to_opt,
        #                        x,
        #                        self.profile,
        #                        p0=guess_arg)
        #
        # fitting = func_to_opt(x, *popt)
        # indicators = calc_mu_sigma(x, *popt)
        # self.integrate(*indicators)https://d2mvzyuse3lwjc.cloudfront.net/images/WikiWeb/Peak_Analysis/PeakAnalyzer_005.png?v=2993

        # It's plotting time ###########################################################################################
        #
        # plt.figure()
        # plt.plot(self.profile_old,
        #           label="Measured profile", linewidth=2)
        # plt.plot(self.baseline_values, "-k", label="Measured profile",
        #           linewidth=2)
        # plt.show()
        # plt.plot(peaks_loc * 50, "-k", label="Measured profile",
        #          linewidth=2)
        # plt.plot(self.profile, "--k", label="Measured profile",
        #          linewidth=2)
        #
        # plt.subplot(3, 5, 14)
        # plt.plot(fitting + self.baseline_values, ":k",
        #          label="Fitted profile", linewidth=2)
        # plt.tick_params(
        #     axis='x',  # changes apply to the x-axis
        #     which='both',  # both major and minor ticks are affected
        #     bottom='off',  # ticks along the bottom edge are off
        #     top='off',  # ticks along the top edge are off
        #     labelbottom='off')
        # plt.ylim([0, 100])
        # plt.ylabel("Signal ($\%$)")
        #plt.legend()

        #plt.plot([indicators[0], indicators[2], indicators[0] - 3 * indicators[1], indicators[0] + 3 * indicators[1]],
         #       [fitting[int(indicators[0])] + self.baseline_values[int(indicators[0])],
          #       fitting[int(indicators[2])] + self.baseline_values[int(indicators[2])],
           #      fitting[int(indicators[0] - 3 * indicators[1])] + self.baseline_values[
            #     fitting[int(indicators[0] + 3 * indicators[1])] + self.baseline_values[
             #        int(indicators[0] + indicators[1] * 3)], 'ro')]


def strip_back(image):
    roi = cv2.imread('/home/alicia/ProjectQR/qrdecode/scripts/images/lila2.png')
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    target = image
    plt.figure()
    plt.imshow(target)
    plt.figure()
    hsvt = cv2.cvtColor(target, cv2.COLOR_BGR2HSV)
    # calculating object histogram
    roihist = cv2.calcHist([hsv], [0, 1], None, [180, 256], [0, 180, 0, 256])
    # normalize histogram and apply back-projection
    cv2.normalize(roihist, roihist, 0, 255, cv2.NORM_MINMAX)
    cv2.imshow('hist', roihist)
    plt.imshow(hsvt)
    dst = cv2.calcBackProject([hsvt], [0, 1], roihist, [0, 180, 0, 256], 1)
    # Now convolve with circular disc
    disc = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    cv2.filter2D(dst, -1, disc, dst)
    # threshold and binary AND
    ret, thresh = cv2.threshold(dst, 50, 255, 0)
    thresh = cv2.merge((thresh, thresh, thresh))
    res = cv2.bitwise_and(target, thresh)
    res = np.vstack((target, thresh, res))
    cv2.imshow('result', res)
    return res
