from guang.cv.MajorTone import *
import cv2
import matplotlib.pyplot as plt
from skimage import data
from skimage.transform import match_histograms


def histogramMatching():
    reference = data.coffee()
    image = data.chelsea()

    matched = match_histograms(image, reference, multichannel=True)

    fig, (ax1, ax2, ax3) = plt.subplots(nrows=1,
                                        ncols=3,
                                        figsize=(8, 3),
                                        sharex=True,
                                        sharey=True)
    for aa in (ax1, ax2, ax3):
        aa.set_axis_off()

    ax1.imshow(image)
    ax1.set_title('Source')
    ax2.imshow(reference)
    ax2.set_title('Reference')
    ax3.imshow(matched)
    ax3.set_title('Matched')

    plt.tight_layout()
    plt.show()


def generateBackGround(img, show_plot=1):
    '''
    input: BGR image array
    return: a background image ,has the same shape with input image.
    '''
    img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    h, w, _ = img.shape
    feature = cv2.resize(img, dsize=(40, 40))
    color1 = MajorTones.getCenters(feature, n_cls=2, m_cls=1)[0][0]
    #     color2 = MajorTones.getCentroid(feature, n_cls = 2)

    color = color1.astype('uint8')
    #     color = color2.astype('uint8')

    if show_plot:
        plt.figure()
        plt.subplot(121)
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_HSV2RGB))
        plt.subplot(122)
        plt.imshow(
            cv2.cvtColor(color.reshape(1, 1, 3),
                         cv2.COLOR_HSV2RGB).reshape(1, 1, 3))
        plt.show()
    H, W = 55, 55
    gener_backg = np.ones((H, W, 3)).astype('uint8')
    gener_backg[:, :, 0] = gener_backg[:, :, 0] * color[0]  # H
    gener_backg[:, :, 1] = gener_backg[:, :, 1] * color[1]  # S
    gener_backg[:, :, 2] = gener_backg[:, :, 2] + 250  # V

    if show_plot:
        plt.figure()
        plt.imshow(cv2.cvtColor(gener_backg, cv2.COLOR_HSV2RGB))
        plt.show()

    i = 1
    c_max = gener_backg[i].max()
    R2 = H**2 + W**2
    for x in range(W):
        for y in range(H):
            r2 = (x - W / 2)**2 + (y - H / 2)**2  # 圆心转移到中心
            c = c_max * r2 // R2
            gener_backg[y, x, i] = c

    gener_backg = cv2.cvtColor(gener_backg, cv2.COLOR_HSV2RGB)
    gener_backg = cv2.resize(gener_backg, (h, w),
                             interpolation=cv2.INTER_NEAREST)
    if show_plot:
        plt.figure()
        plt.imshow(gener_backg)
        plt.show()

    return gener_backg


if __name__ == "__main__":

    def test_generateBackGround():
        img = cv2.imread("pic/hue1.png")
        generateBackGround(img, show_plot=1)

    # test_generateBackGround()
    histogramMatching()
