#!/usr/bin/python
import cv2
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.cm as cm
import random
import sys
# from PIL import Image
PATH = 'lab_6_images'

def BGR2RGB(img):
	res = np.zeros(img.shape , np.uint8)
	res[:,:,0]=img[:,:,2]
	res[:,:,1]=img[:,:,1]
	res[:,:,2]=img[:,:,0]
	return res

def process(img1):
	QY=np.array([[16,11,10,16,24,40,51,61]
		,[12,12,14,19,26,48,60,55]
		,[14,13,16,24,40,57,69,56]
		,[14,17,22,29,51,87,80,62]
		,[18,22,37,56,68,109,103,77]
		,[24,35,55,64,81,104,113,92]
		,[49,64,78,87,103,121,120,101]
		,[72,92,95,98,112,100,103,99]])

	QC=np.array([[17,18,24,47,99,99,99,99]
		,[18,21,26,66,99,99,99,99]
		,[24,26,56,99,99,99,99,99]
		,[47,66,99,99,99,99,99,99]
		,[99,99,99,99,99,99,99,99]
		,[99,99,99,99,99,99,99,99]
		,[99,99,99,99,99,99,99,99]
		,[99,99,99,99,99,99,99,99]])

	B=8 # blocksize

	h,w=np.array(img1.shape[:2])/B * B
	img1=img1[:h,:w]

	img2 = BGR2RGB(img1)
	# plt.imshow(img2)

	#convert to YCrCb
	transcol=cv2.cvtColor(img1, cv2.cv.CV_BGR2YCrCb)

	#sampling
	SSV=2
	SSH=2
	crs=cv2.boxFilter(transcol[:,:,1],ddepth=-1,ksize=(2,2))
	cbs=cv2.boxFilter(transcol[:,:,2],ddepth=-1,ksize=(2,2))

	imSub=[transcol[:,:,0],crs[::SSV,::SSH],cbs[::SSV,::SSH]]

	Qf=1
	if Qf < 50 and Qf > 1:
		scale = np.floor(5000/Qf)
	elif Qf < 100:
		scale = 200-2*Qf
	else:
		raise Exception('wat?')

	scale=scale/100.0


	Q=[QY*scale,QC*scale,QC*scale]
	TransAll=[]
	TransAllFiltred=[]
	TransAllQuant=[]
	TransAllQuantFiltred=[]


	for idx,channel in enumerate(imSub):
		channelrows=channel.shape[0]
		channelcols=channel.shape[1]
		Trans = np.zeros((channelrows,channelcols), np.float32)
		TransQuant = np.zeros((channelrows,channelcols), np.float32)
		blocksV=channelrows/B
		blocksH=channelcols/B
		vis0 = np.zeros((channelrows,channelcols), np.float32)
		vis0[:channelrows, :channelcols] = channel
		vis0=vis0-128
		for row in xrange(blocksV):
				for col in xrange(blocksH):
						currentblock = cv2.dct(vis0[row*B:(row+1)*B,col*B:(col+1)*B])
						Trans[row*B:(row+1)*B,col*B:(col+1)*B]=currentblock
						TransQuant[row*B:(row+1)*B,col*B:(col+1)*B]=np.round(currentblock/Q[idx])
		TransAll.append(Trans)
		TransAllQuant.append(TransQuant)

		tq1 = TransQuant[TransQuant!=0]
		tq2 = tq1[tq1!=1]
		tq3 = tq2[tq2!=-1]


		TransAllQuantFiltred.append(tq3)
		TransAllFiltred.append(Trans[(Trans!=0)&(Trans!=1)&(Trans!=-1)])

	return TransAllQuantFiltred

def test(path, bins, cfactor):
	img = cv2.imread(path, cv2.CV_LOAD_IMAGE_UNCHANGED)
	img2 = img[cfactor:len(img)-cfactor, cfactor:len(img[0])-cfactor]
		# img2 = cv2.GaussianBlur(img2,(3,3),0)
		# img = cv2.GaussianBlur(img,(3,3),0)
	t1 = process(img)
	t2 = process(img2)

	plt.figure(path+'0')

	plt.hist(t1[0], bins,  alpha=0.75)
	plt.hist(t2[0], bins,  alpha=0.75)
	# plt.figure(path+'1')
	# plt.hist(t1[1], bins,  alpha=0.75)
	# plt.hist(t2[1], bins,  alpha=0.75)

	# plt.figure(path+'2')
	# plt.hist(t1[2], bins,  alpha=0.75)
	# plt.hist(t2[2], bins,  alpha=0.75)
	plt.show()

def main(filename):
	bins = 240
	cfactor = 5
	test(filename, bins, cfactor)

if __name__ == '__main__':
	print(sys.argv)
	main(sys.argv[1])