import time
import cv2
import numpy as np
from matplotlib import pyplot as plt
from collections import Counter


import matplotlib
matplotlib.rcParams['savefig.dpi'] = 4 * matplotlib.rcParams['savefig.dpi']

img = cv2.imread('../tessdata/1.png', cv2.CV_LOAD_IMAGE_GRAYSCALE)[245:727+245, 90:1198+90]

import math
def q(x):
    return math.sqrt(x*x*x)

def findTextColor(img):
    tmp = np.copy(img[100:200,100:200])
    tmp[tmp < 100] = 0

    contours, _ = cv2.findContours(tmp, cv2.RETR_LIST, cv2.RETR_LIST )

    colors = []
    for c in contours:
        region = tmp[c]
        region = region[region > 100]
        colors.append(max(region.flatten()))
        #colors.append(np.sum(region) / len(region))
    return max(colors)
#    return np.sum(colors) / len(colors)

print 'text color:',findTextColor(img)

def invertSelectionLine(img):
    def detectSelectionLine(img):
        for i in range(img.shape[0]):
            if np.sum(img[i,:]) / img.shape[1] > 100:
                for j in range(i, img.shape[0]):
                    if np.sum(img[j,:]) / img.shape[0] <= 100:
                        return i, j
                return i, img.shape[0]
    firstRow, lastRow = detectSelectionLine(img)
    textColor = findTextColor(img)
    inv = img[firstRow:lastRow,:].astype('int16')
    inv = 255 - inv

    globalAvg = np.average(img)
    areaAvg = np.average(inv)
    inv -= areaAvg - globalAvg
    inv[inv<0]=0
    img[firstRow:lastRow,:] = inv.astype('uint8')

def normalizeColors(img):
    tmp = (img.astype('int16') * 1.2)
    maxColor = max(tmp.flatten())
    maxColor -= 255
    tmp -= maxColor
    tmp[tmp<0] = 0
    tmp *= 3
    tmp[tmp>255] = 255
    return img


def contrast(img):
    THRESHOLD = 80
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if img[i,j] > THRESHOLD:
                img[i,j] += min(255 - img[i,j], q(img[i,j]-THRESHOLD))
            else:
                img[i,j] -= min(img[i,j], q(THRESHOLD - img[i,j]))

def contrast(img, mul=2, shift=0):
    tmp = np.copy(img).astype(float) * mul
    avg = np.sum(tmp.flatten()) / len(tmp.flatten()) + shift
    tmp -= avg
    tmp[tmp < 0] = 0
    tmp[tmp > 255] = 255
    return tmp.astype('uint8')

def detectVerticalLines(img):
    for col in range(2, img.shape[1]):
        diff = np.sum(abs(img[:,col] - img[:,col-1])) / img.shape[0]
        if diff > 200:
            for col2 in range(col+1, img.shape[1]):
                diff2 = np.sum(abs(img[:,col2] - img[:,col2-1])) / img.shape[0]
                if diff2 > 200:
                    width = col2 - col
                    hw = width / 2
                    if hw == 0: hw = 1

                    mid = (col2 + col) / 2

                    img[:,col:mid] = img[:,col-hw-1:col-1]
                    img[:,(col+col2)/2:col2] = img[:,col2+1:col2+1+hw]
                    return

#for i in range(3):
#    detectVerticalLines(img)

def removeVerticalLines(img):
    minrows = [r[1] for r in sorted([(np.sum(img[row,:]),row) for row in range(img.shape[0])])[:50]]
    cols = []
    for col in range(100,img.shape[1]):
        avg = np.sum(img[minrows,col]) / len(minrows)
        cols.append(avg)

    freq = np.fft.fft(cols)
    freq[100:] = 0
    diff = abs(cols - np.fft.ifft(freq))
    maxDiff = max(diff)
    for i in range(len(diff)):
        if diff[i] > maxDiff *0.20:
            img[:,i+100] = 0#np.sum(img.flatten()) / len(img.flatten())

def removeHorizontalLines(img):
    mincols = [r[1] for r in sorted([(np.sum(img[:,col]),col) for col in range(img.shape[1])])[:50]]
    rows = []
    for row in range(100,img.shape[0]):
        avg = np.sum(img[row, mincols]) / len(mincols)
        rows.append(avg)

    freq = np.fft.fft(rows)
    freq[100:] = 0
    diff = abs(rows - np.fft.ifft(freq))
    maxDiff = max(diff)
    for i in range(len(diff)):
        if diff[i] > maxDiff *0.20:
            img[i+100,:] = 0#np.sum(img.flatten()) / len(img.flatten())




#img = contrast(img, 2, 40)
#img = contrast(img, 1.5, 20)
#img = contrast(img, 1.5, 10)
#img[img<10] = 0
#img[img>0] = 255

def fftRows(img):
    for i in range(img.shape[0]):
        row = img[i]
        dark = row[row < 100]
        fimg = np.fft.fft(dark)
        fimg[0] = 0
        row[row < 100] = abs(np.fft.ifft(fimg))


def cleanBackgroundFFT(img):
    print 'fft start'
    start = time.time()
    dark = img[img < 100]
    chunkSize = img.shape[1] * 10
    chunkCount = int(math.ceil(float(len(dark))/chunkSize)*chunkSize)
    for chunk in range(0, chunkCount, chunkSize):
        fimg = np.fft.fft(dark[chunk:chunk+chunkSize])
        fimg[:10] = 0
        dark[chunk:chunk+chunkSize] = abs(np.fft.ifft(fimg))
    img[img<100] = dark
    end = time.time()
    print 'fft end'
    print 'fft took', end-start


#img[row,:] = abs(np.fft.ifft(fimg))

#fimg = np.fft.fft2(img)
#fimg[0] /= 1000
#fimg[1] /= 1000
#img = abs(np.fft.ifft2(fimg))


invertSelectionLine(img)
removeHorizontalLines(img)
cleanBackgroundFFT(img)
removeVerticalLines(img)
img = contrast(img, 3, 20)
img = contrast(img, 1.5, 10)

cv2.imwrite('q.tiff', img)
#plt.imshow(img, 'gray')