import os

from PIL import Image
from matplotlib import pyplot as plt

import numpy as np


def img_to_grey_mat(image):
    greys = []
    for j in xrange(image.height):
        for i in xrange(image.width):
            rgb = image.getpixel((i, j))
            if isinstance(rgb, tuple):
                val = __rgb2grey(rgb[0], rgb[1], rgb[2])
            else:
                val = rgb
            greys.append(val)

    mat = np.array(greys, np.int16)
    mat.resize(image.height, image.width)
    return mat


def img_to_binary_mat(image, minumum_brightness=128,
                      brighter=True, shape=None):
    """
    binarize a PIL Image object to a binarized numpy array, 0, 1
    """
    if not shape:
        shape = image.height, image.width
    mat = img_to_grey_mat(image)
    binarization_mat(mat, minumum_brightness, brighter)
    mat = patch_2_regular_shape(mat, shape)
    return mat


def __rgb2grey(r, g, b):
    return (r * 19595 + g * 38469 + b * 7472) >> 16


def binarization_mat(mat, valve=128, brighter=True):
    """
    binarization_mat
    binarize a grey numpy array to 0, 1
    """
    mat[mat < valve] = 0
    mat[mat >= valve] = 1
    if not brighter:
        mat[mat == 1] = 2
        mat[mat == 0] = 1
        mat[mat == 2] = 0
    return mat


def patch_2_regular_shape(mat, shape):
    """
    as name indicates
    mat is a binarized numpy array
    """
    y_start = (shape[0] - mat.shape[0]) // 2
    x_start = (shape[1] - mat.shape[1]) // 2
    m = np.zeros(shape, dtype=np.int32)
    m[y_start:y_start + mat.shape[0], x_start:x_start + mat.shape[1]] = mat
    return m


def show_mat(mat):
    """
    plot mat
    """
    ax = plt.subplot()
    plt.subplot()
    plt.imshow(mat, cmap='gray', interpolation='none')
    ax.axis('off')
    plt.show()


def show_mats(mats, line_width=10):
    row_nums = int(np.ceil(len(mats) * 1.0 / line_width))
    n = 0
    for mat in mats:
        n += 1
        ax = plt.subplot(row_nums, line_width, n)
        ax.axis('off')
        ax.imshow(mat, cmap='gray', interpolation='none')
    plt.show()


def save_mat_as_img(mat, fname):
    plt._imsave(arr=mat, fname=fname, cmap="gray")


def save_mats_as_imgs(mats, dname="./", fname_prefix="Img_", start_index=1):
    if not os.path.exists(dname):
        os.mkdir(dname)

    index = start_index
    for mat in mats:
        f_name = dname + "/" + fname_prefix + str(index) + ".jpg"
        save_mat_as_img(mat, f_name)
        index += 1


def remove_singular_noise(mat, pat):
    """
    remove singular noise
    if a noise looks like:
    [0,0,0]
    [0,1,0]
    [0,0,0]
    it's pattern will be:
    [1,1,1]
    [1,2,1]
    [1,1,1]
    """
    H, W = mat.shape

    # fine me
    me_xy = np.where(pat > 1)
    me_x = me_xy[0][0]
    me_y = me_xy[1][0]

    neighbours_xys = np.where(pat == 1)

    # find neighbours coordinate diffs
    xx = neighbours_xys[0] - me_x
    yy = neighbours_xys[1] - me_y

    mat[[0, -1], :] = 0
    mat[:, [0, -1]] = 0
    for h in xrange(1, H - 1):
        for k in xrange(1, W - 1):
            if mat[h, k] > 0:
                vals = []
                for i_xy in xrange(len(xx)):
                    val = 0
                    h1, k1 = h + yy[i_xy], k + xx[i_xy]
                    if h1 < 0 or h1 >= H or k1 < 0 or k1 >= W:
                        val = 0
                    else:
                        val = mat[h1, k1]
                    vals.append(val)
                if not np.any(np.array(vals)):
                    mat[h, k] = 0


def center_patch_to_regular_size(mat, regular_shape):
    M, N = mat.shape
    xproject = np.sum(mat, 0)
    xsum = np.dot(np.arange(N, dtype=np.int16), xproject)
    yproject = np.sum(mat, 1)
    ysum = np.dot(np.arange(M, dtype=np.int16), yproject)

    xm = 1.0 * xsum / sum(xproject)
    ym = 1.0 * ysum / sum(yproject)

    cm, cn = int(round(ym)), int(round(xm))

    m, n = regular_shape

    y1, y2 = 0, m - 1
    halfy = m // 2
    if m % 2:  # odd
        Y1 = cm - halfy
        Y2 = cm + halfy
    else:  # even
        Y1 = cm - halfy + 1
        Y2 = cm + halfy
    if Y1 < 0:
        y1 -= Y1
        Y1 = 0
    if Y2 > M - 1:
        d = Y2 - M + 1
        y2 -= d
        Y2 = M - 1

    x1, x2 = 0, n - 1
    halfx = n // 2
    if n % 2:  # odd
        X1 = cn - halfx
        X2 = cn + halfx
    else:  # even
        X1 = cn - halfx + 1
        X2 = cn + halfx
    if X1 < 0:
        x1 -= X1
        X1 = 0
    if X2 > N - 1:
        d = X2 - N + 1
        x2 -= d
        X2 = N - 1

    res = np.zeros(regular_shape, dtype=np.int16)
    res[y1:y2 + 1, x1:x2 + 1] = mat[Y1:Y2 + 1, X1:X2 + 1]
    return res


if __name__ == '__main__':
    img = Image.open('F:/shanghai_si_captchas/1.jpg')
    mat = img_to_binary_mat(img, minumum_brightness=180, brighter=False)
    print mat
    print np.sum(mat.reshape(mat.shape[0] * mat.shape[1]))
    print mat.shape
    m1 = center_patch_to_regular_size(mat, regular_shape=(20, 20))
    show_mat(m1)