import numpy as np 

# calculate output shape 
def calc_shape(input_, kernal_, stride = 1, padding = 0):
	H_out = 1+(input_.shape[0]+2*padding-kernal_.shape[0])/stride
	W_out = 1+(input_.shape[1]+2*padding-kernal_.shape[1])/stride
	return H_out,W_out

def conv2d_naive(input_, kernal_, stride =1, padding = 0):
	# calculate Convolution param 
	out = None
	H,W = input_.shape
	x_pad = np.zeros((H+2*padding,W+2*padding))
	x_pad[padding:padding+H,padding:padding+W] = input_
	H_out,W_out = calc_shape(input_, kernal_, stride, padding)
	out = np.zeros((H_out,W_out))
	# Convolution
	for m in xrange(H_out):
		for n in xrange(W_out):
			out[m,n] = np.sum(x_pad[m*stride:m*stride+kernal_.shape[0],n*stride:n*stride+kernal_.shape[1]] *  kernal_)
	return out

import numpy as np 

# calculate output shape 
def calc_shape(input_, kernal_, stride = 1, padding = 0):
	H_out = 1+(input_.shape[0]+2*padding-kernal_.shape[0])/stride
	W_out = 1+(input_.shape[1]+2*padding-kernal_.shape[1])/stride
	return H_out,W_out

def im2col(input_, kernal_):
	# calculate param of col matrix
	X_ = input_.shape[0]-kernal_.shape[0]+1
	Y_ = input_.shape[1]-kernal_.shape[1]+1
	output_col = np.empty((kernal_.shape[0]*kernal_.shape[1], X_*Y_))
	# im2col
	for i in range(Y_):
		for j in range(X_):
			output_col[:,i*Y_+j] = input_[j:j+kernal_.shape[0],i:i+kernal_.shape[1]].ravel(order='F')
	return 大专栏  卷积神经网络基本概念（卷积篇 01） output_col

def col2im(input_, kernal_ ,col_matrix):
	output_ = np.zeros(input_.shape) 
	weight_ = np.zeros(input_.shape)
	col = 0
	X_ = input_.shape[0] - kernal_.shape[0] + 1
	Y_ = input_.shape[1] - kernal_.shape[1] + 1
	for i in range(Y_):
		for j in range(X_):
			output_[j:j+kernal_.shape[0],i:i+kernal_.shape[1]]+= col_matrix[:,col].reshape(kernal_.shape, order='F')
			weight_[j:j+kernal_.shape[0],i:i+kernal_.shape[1]] += np.ones(kernal_.shape)
			col+=1
	return output_/weight_

def conv2d_im2col(input_, kernal_, stride =1, padding = 0):
	# stride must = 1 in this code 
	# this code is only show how im2col work
	# if you want stride = other number , change the np.dot
	H,W = input_.shape
	x_pad = np.zeros((H+2*padding ,W+2*padding))
	x_pad[padding:padding+H,padding:padding+W] = input_

	# im2col processing and dot col and kernal
	output_col = im2col(x_pad, kernal_)
	col_matrix = np.dot(kernal_.reshape(kernal_.shape[0]*kernal_.shape[1],order='F'),output_col)
	
	# reshape to the final status
	H_out,W_out = calc_shape(input_, kernal_, stride, padding)
	out = np.zeros((H_out,W_out))
	return col_matrix.reshape(out.shape[0],out.shape[1],order='F')

def FFT_convolve(input_,kernal_):
	# calculate the FFT trans size
	T_size = np.array(input_.shape)+np.array(kernal_.shape)-1
	# ceil doc : https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.ceil.html
	F_size = 2**np.ceil(np.log2(T_size)).astype(int)
	F_slice = tuple([slice(0, int(item)) for item in F_size])

	new_input_ = np.fft.fft2(input_, F_size)
	new_kernal_ = np.fft.fft2(kernal_, F_size)
	output_ = np.fft.ifft2(new_input_*new_kernal_)[F_slice].copy()
	# output is a expand matrix which bigger than the result you suppose
	return np.array(output_.real, np.int32)

def demo_1():
	input_ = np.diag([1, 10, 3, 5, 1])
	kernal_ = np.diag([1, 2, 1])
	return conv2d_naive(input_, kernal_, stride =1, padding = 0)

def demo_2():
	input_ = np.diag([1, 1, 1, 1, 1])
	kernal_ = np.diag([1, 1, 1])
	# stride must equal 1 in my demo code
	return conv2d_im2col(input_, kernal_, stride =1, padding = 0)

def demo_3():
	input_ = np.diag([1, 10, 3, 5, 1])
	kernal_ = np.diag([1, 2, 1,])
	return FFT_convolve(input_, kernal_)

print demo_1()
print demo_2()
print demo_3()