import numpy as np
from scipy.io import loadmat


class CNN:
    def __init__(self, layer1=2, learning_rate=0.1, iters=10000):
        self.layer1 = layer1  # 第一個卷積層卷積核的個數
        self.iters = iters  # 最大迭代次數
        self.learning_rate = learning_rate  # 學習率
        self.maxindex = []  # 儲存池化區域最大值索引
        self.k = 0  # 池化後矩陣索引
        self.cost = []  # 損失值
        self.parameter = []  # 儲存訓練好的引數

    @staticmethod
    def relu(mat):  
        """定義啟用函式relu,對矩陣mat的所有元素操作一遍"""
        for i in range(mat.shape[0]):
            for j in range(mat.shape[1]):
                mat[i, j] = max(0, mat[i, j])
        return mat
    
    @staticmethod
    def sigmoid(x):
        """定義sigmoid函式"""
        return 1 / (1 + np.exp(x))

    @staticmethod
    def softmax(x):
        """定義softmax函式,x是列向量"""
        return np.exp(x) / np.sum(np.exp(x))

    @staticmethod
    def cal_cost(y_pred, y_ture):
        """
        計算平方損失誤差
        :param y_pred: 預測值k*1
        :param y_ture: 真實值
        :return: 平方誤差
        """
        return np.sum(np.power(y_pred - y_ture, 2)) / 2

    @staticmethod
    def cal_conv(mat_a, mat_b, b=0, step=1, padding=0):
        """
        二維卷積
        對矩陣mat_a填充padding後,使用卷積核大小為F*F的矩陣,以步伐step進行卷積計算,返回卷積後的矩陣
        :param mat_a: 被卷積矩陣h*w
        :param mat_b: 卷積核矩陣f*f
        :param b: 卷積層偏置
        :param step: 步伐S
        :param padding: 對矩陣mat_a外圍的填充0的層數
        :return: 返回卷積後的矩陣
        """
        h, w, f = mat_a.shape[0], mat_a.shape[1], mat_b.shape[0]  # 高,寬,卷積核邊長（卷積核為方陣,且通常邊長為奇數）
        conved_h, conved_w = (h - f + 2 * padding) // step + 1, (w - f + 2 * padding) // step + 1
        conved = np.mat(np.zeros((conved_h, conved_w)))  # 定義卷積後矩陣
        if padding:  # 填充操作
            new_a = np.mat(np.zeros((h + 2 * padding, w + 2 * padding)))
            new_a[padding: -padding, padding: - padding] = mat_a
            mat_a = new_a
        for i in range(conved_h):  # 卷積計算
            for j in range(conved_w):
                conved[i, j] = np.sum(np.multiply(mat_a[i: i + f, j: j + f], mat_b))
        return conved + b

    def pool(self, mat, f=2):
        """
        對矩陣mat,使用邊長為f的矩陣池化,採用最大值池化法,返回池化後的矩陣
        :param mat: 被池化層H*w
        :param f: 過濾器的邊長
        :return: 池化後的矩陣
        """
        h, w = mat.shape[0], mat.shape[1]  # 高,寬
        pooled_h, pooled_w = h // f, w // f
        pooled = np.mat(np.zeros((pooled_h, pooled_w)))
        for i in range(pooled_h):
            for j in range(pooled_w):
                temp = mat[i * f: i * f + f, j * f: j * f + f]
                self.maxindex.append(np.argmax(temp))  # 記錄最大值索引
                pooled[i, j] = np.max(temp)
        return pooled

    def upsample(self, mat, f=2):
        """
        對梯度誤差矩陣A,執行最大池化法的上取樣,還原成輸入矩陣的結構
        :param mat: 梯度誤差矩陣
        :param f: 原池化區域大小
        :return: 上一層的梯度誤差矩陣
        """
        pooled_h, pooled_w = mat.shape
        h, w = pooled_h * f, pooled_w * f
        origin = np.mat(np.zeros((h, w)))  # 定義原始矩陣
        for i in range(pooled_h):
            for j in range(pooled_w):
                temp = origin[i * f: i * f + f, j * f: j * f + f]
                temp[self.maxindex[self.k] // f, self.maxindex[self.k] % f] = mat[i, j]
                origin[i * f: i * f + f, j * f: j * f + f] = temp
                self.k += 1
        return origin

    @staticmethod
    def creat_conv(n_conv, f, dim3=1):
        """
        建立卷積層,初始化引數
        :param n_conv: 卷積核個數
        :param f: 卷積核邊長
        :param dim3: 輸入張量第三維的數值,預設為1,表示輸入是一個矩陣,若為2,則表示輸入了2個矩陣,以此類推
        :return: 卷積層的初始引數
        """
        weight = []
        for _ in range(n_conv):
            temp = []
            for _ in range(dim3):
                temp.append(np.mat((2 * np.random.rand(f, f) - np.ones((f, f))) / (np.sqrt(n_conv) * 100)))  # 儘量使引數值小
            weight.append(temp)
        b = np.mat((2 * np.random.rand(n_conv, 1) - np.ones((n_conv, 1))) / (np.sqrt(n_conv) * 100))
        return weight, b

    @staticmethod
    def creat_fc(n_neural, n_put):
        """
        建立全連線層,初始化引數
        :param n_put: 上一層層神經元個數
        :param n_neural: 全連線層神經元個數
        :return: 全連線層的引數
        """
        weight = np.mat((2 * np.random.rand(n_neural, n_put) - np.ones((n_neural, n_put))) / (np.sqrt(n_neural) * 100))
        b = np.mat((2 * np.random.rand(n_neural, 1) - np.ones((n_neural, 1))) / (np.sqrt(n_neural) * 100))
        return weight, b

    def training(self, features, target):
        """
        根據CNN演算法訓練模型,使用的結構為
        input->conv->pool->fc->output
        本例只考慮通道為1的灰度影象
        :param features: 特徵集m*n,m為樣本個數,n為圖片畫素總數,如圖片為28*28,則n=784,訓練時,要把一維資料重新轉換成28*28的矩陣
        :param target: 標籤集m*k,k為類別數量,要求y已進行過獨熱編碼
        :return: 模型引數
        """
        m, k = features.shape[0], target.shape[1]
        features = np.mat(features)
        target = np.mat(target)
        weight2, b2 = self.creat_conv(self.layer1, 5, dim3=1)
        weight4, b4 = self.creat_fc(100, 288)
        weight5, b5 = self.creat_fc(k, 100)

        for index in range(self.iters):
            print(index)
            if index == m:
                break
            train_x = features[index, :]
            y_true = target[index, :].T
            train_x = train_x.reshape(28, 28)

            a1 = [train_x]
            a2 = []
            for i in range(self.layer1):  # 第二層卷積核個數
                temp = np.mat(np.zeros((24, 24)))
                for j in range(1):  # dim3=1
                    temp += self.cal_conv(a1[j], weight2[i][j], b=b2[i, 0])
                a2.append(self.relu(temp))
            a3 = []
            for i in range(self.layer1):
                a3.append(self.pool(a2[i]))
            a3flat = np.mat(np.zeros((1, self.layer1 * 12 * 12)))
            for i in range(self.layer1):
                a3flat[0, 144 * i: 144 * i + 144] = a3[i].flatten()
            a3 = a3flat.T
            a4 = self.sigmoid(np.dot(weight4, a3) + b4)
            a5 = self.softmax(np.dot(weight5, a4) + b5)
            self.cost.append(self.cal_cost(a5, y_true))

            delta5 = np.multiply(a5, np.multiply(a5 - y_true, 1 - a5))  # 計算第五層delta和梯度
            grad_weight5 = np.dot(delta5, a4.T)
            grad_b5 = delta5

            delta4 = np.multiply(np.dot(weight5.T, delta5), np.multiply(a4, 1 - a4))  # 計算第四層delta和梯度
            grad_weight4 = np.dot(delta4, a3.T)
            grad_b4 = delta4

            delta3 = np.dot(weight4.T, delta4)[::-1]  # 計算第三層delta和梯度
            delta3_ = []
            for i in range(self.layer1):
                temp = delta3[i * 144: i * 144 + 144, 0].reshape(12, 12)
                delta3_.append(temp)
            delta3 = delta3_

            delta2 = []  # 計算第二層delta和梯度
            for i in range(self.layer1):
                delta2.append(self.upsample(delta3[i]))

            grad_weight2 = []
            grad_b2 = np.mat(np.zeros((self.layer1, 1)))
            for i in range(self.layer1):
                temp = []
                for j in range(1):
                    temp.append(self.cal_conv(a1[j], delta2[i]))
                grad_b2[i, 0] = np.sum(delta2[i])
                grad_weight2.append(temp)

            # 更新引數值
            weight5 -= self.learning_rate * grad_weight5
            b5 -= self.learning_rate * grad_b5
            weight4 -= self.learning_rate * grad_weight4
            b4 -= self.learning_rate * grad_b4
            for i in range(self.layer1):
                for j in range(1):
                    weight2[i][j] -= self.learning_rate * grad_weight2[i][j]
            b2 -= self.learning_rate * grad_b2

            self.maxindex = []  # 下一輪迭代器重置
            self.k = 0

        self.parameter.extend([weight2, b2, weight4, b4, weight5, b5])  # 儲存引數
        return

    def predict(self, features):
        features = np.mat(features)
        preiction = []
        m = features.shape[0]
        for index in range(m):
            x = features[index, :]
            x = x.reshape(28, 28)
            weight2, b2, weight4, b4, weight5, b5 = self.parameter
            # 前向傳播求輸出
            a1 = [x]
            a2 = []
            for i in range(self.layer1):  # 第二層卷積核個數
                temp = np.mat(np.zeros((24, 24)))
                for j in range(1):  # dim3=1
                    temp += self.cal_conv(a1[j], weight2[i][j], b=b2[i, 0])
                a2.append(self.relu(temp))
            a3 = []
            for i in range(self.layer1):
                a3.append(self.pool(a2[i]))
            a3flat = np.mat(np.zeros((1, self.layer1 * 12 * 12)))
            for i in range(self.layer1):
                a3flat[0, 144 * i: 144 * i + 144] = a3[i].flatten()
            a3 = a3flat.T
            a4 = self.sigmoid(np.dot(weight4, a3) + b4)
            a5 = self.softmax(np.dot(weight5, a4) + b5)

            # 本例中最大值索引即為對應的數字
            max_index = np.argmax(a5)
            preiction.append(max_index)

        return preiction


def test():
    """使用手寫數字集進行測試,效果不理想,不知道問題出在哪"""
    dataset = loadmat('data/mnist_all')
    train_features = np.zeros((1, 784))
    train_target = np.zeros((1, 10))
    valid_features = np.zeros((1, 784))
    valid_target = []
    for i in range(10):  # 每個數字獲取1000個訓練資料和50個驗證資料
        temp_train = dataset['train%d' % i]
        temp_valid = dataset['test%d' % i]
        rand_index = np.random.choice(temp_train.shape[0], 1000)  # 隨機選取1000個訓練樣本
        valid_index = np.random.choice(temp_valid.shape[0], 50)
        temp_train = temp_train[rand_index, :]
        temp_valid = temp_valid[valid_index, :]
        train_features = np.concatenate((train_features, temp_train))
        valid_features = np.concatenate((valid_features, temp_valid))
        target = np.zeros((1000, 10))
        target[:, i] = 1
        train_target = np.concatenate((train_target, target))
        valid_target.extend([i] * 50)

    train_features = train_features[1:, :]
    train_target = train_target[1:, :]
    valid_features = valid_features[1:, :]
    train_data = np.concatenate((train_features, train_target), axis=1)
    np.random.shuffle(train_data)  # 洗牌
    train_features = train_data[:, : 784]
    train_target = train_data[:, 784:]

    cnn = CNN(learning_rate=0.1, iters=1000)
    cnn.training(train_features, train_target)
    prediction = cnn.predict(valid_features)
    correct = [1 if a == b else 0 for a, b in zip(prediction, valid_target)]
    print(correct.count(1) / len(correct))
    cost = cnn.cost
    print(cost[::100])


test()