import os
import sys
import numpy as np
import numpy

# 
def conv_(img, conv_filter, stride = 1):
    """
    img: wxh 二维图像
    conv_filter: kxk 二维卷积核(eg. 3x3)
    """
    filter_size = conv_filter.shape[1]
    out_w = np.uint16((img.shape[0]-conv_filter.shape[0]) / stride + 1)
    out_h = np.uint16((img.shape[1]-conv_filter.shape[1]) / stride + 1)
    result = np.zeros((out_w, out_h))

    for r in range(out_w):
        for c in range(out_h):
            cur_img_region = img[r*stride:r*stride+filter_size, c*stride:c*stride+filter_size]
            cur_result = cur_img_region * conv_filter
            conv_sum = np.sum(cur_result)
            result[r, c] = conv_sum

    return result


def conv(img, conv_filter, padding=0, stride=1):
    if len(img.shape) != len(conv_filter.shape) - 1:
        print("Error: Number of dimensions in conv filter and image do not match.")  
        exit()

    out_w = np.int((img.shape[0] - conv_filter.shape[1] + 2*padding) / stride + 1)
    out_h = np.int((img.shape[1] - conv_filter.shape[2] + 2*padding) / stride + 1)
    out_c = conv_filter.shape[0]
    feature_maps = np.zeros((out_w, out_h, out_c))

    if padding>0:# 先对img进行padding
        new_img = np.zeros((img.shape[0] + 2*padding, img.shape[1] + 2*padding,img.shape[2]))
        new_img[padding:-padding, padding:-padding, :] = img.copy()
        img = new_img

    for filter_num in range(conv_filter.shape[0]):
        cur_filter = conv_filter[filter_num, :]

        if len(cur_filter.shape) > 2:
            conv_map = conv_(img[:, :, 0], cur_filter[:, :, 0])
            for ch_num in range(1, cur_filter.shape[-1]):
                conv_map = conv_map + conv_(img[:, :, ch_num], cur_filter[:, :, ch_num])
        else:
            conv_map = conv_(img, cur_filter)
        
        feature_maps[:, :, filter_num] = conv_map # Holding feature map with the current filter.
    return feature_maps


def pooling(feature_map, kernel_size=2, stride=2):
    out_w = np.int((feature_map.shape[0] - kernel_size) / stride + 1)
    out_h = np.int((feature_map.shape[1] - kernel_size) / stride + 1)
    out_c = feature_map.shape[2]
    pool_out = np.zeros((out_w, out_h, out_c))
    for ch_num in range(out_c):
        for r in range(out_w):
            for c in range(out_h):
                pool_out[r, c, ch_num] = np.max(feature_map[r*stride:r*stride+kernel_size,
                                                            c*stride:c*stride+kernel_size])
    return pool_out


def relu(feature_map):
    for ch_num in range(feature_map.shape[2]):
        for r in range(feature_map.shape[0]):
            for c in range(feature_map.shape[1]):
                feature_map[r,c,ch_num] = max(0, feature_map[r,c,ch_num])
    return feature_map


def conv_Test():
    img = np.random.randint(1, 10, (8,8))
    print(img)
    conv_filter = np.random.randint(1, 10, (5,5))
    print(conv_filter)

    r1 = conv_(img, conv_filter)
    print(r1.shape)
    print(r1)
    # r2 = conv2_(img, conv_filter)
    # print(r2.shape)
    # print(r2)

def convTest():
    img = np.random.randint(1, 10, (8,8,3))
    print(img.shape)
    conv_filter = np.random.randint(1, 10, (6,5,5,3))
    print(conv_filter.shape)
    r1 = conv(img, conv_filter)
    print(r1)
    print(r1.shape)

    import NumPyCNN
    r2 = NumPyCNN.conv(img, conv_filter)
    print(r2)
    print(r2.shape)    

def poolingTest():
    img = np.random.randint(1, 100, (6,6,1))
    print(img[:,:,0])
    print(img.shape)
    pool_out = pooling(img)
    print(pool_out[:,:,0])
    print(pool_out.shape)

# convTest()
poolingTest()


# https://www.linkedin.com/pulse/building-convolutional-neural-network-using-numpy-from-ahmed-gad/