import time
import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm

os.chdir("D:/下載/train/train")

# 讀取圖片的數量
n = 1000
weight = 64
height = 64


def init():
    imgs = np.zeros((n, weight, height, 3), dtype=np.uint8)
    imgs_lable = np.zeros((n, 1), dtype=np.uint8)
    plt.figure(figsize=(16, 8))
    for i in tqdm(range(n)):
        if(i % 2 == 0):
            imgs[i] = cv2.resize(cv2.imread('cat.%d.jpg' %
                                 i), (weight, height))
            imgs_lable[i] = 1
        else:
            imgs[i] = cv2.resize(cv2.imread('dog.%d.jpg' %
                                 i), (weight, height))
        if(i < 10):
            plt.subplot(2, 5, i+1)
            plt.imshow(imgs[i])
            plt.title('imgs_lable:'+str(imgs_lable[i][0]))
    train = imgs.reshape(imgs.shape[0], -1).T/255.
    train_lable = imgs_lable.T
    print("imgs_shape:", imgs.shape, 'train.shape:',
          train.shape, 'train_lable.shape:', train_lable.shape)
    # plt.show()
    return train, train_lable


def initWB(dim, annNum=1):
    w = np.random.randn(dim, annNum)*0.01
    b = 0
    return w, b


def sigmoid(z):
    a = 1.0/(1+np.exp(-z))
    return a


def relu(z):
    a = np.maximum(z, 0)
    return a


def Hidden_layer_Forward_propagation(w, b, x):
    Z = np.dot(w.T, x) + b
    A = relu(Z)
    return Z, A


# W2輸出層的權重，dz2輸出層z的偏導,dg輸出層啟用函式的偏導
def Hidden_layer_Backward_propagation(x, Z, w2, dz2, m):
    # relu的偏導是大於等於0是1，不然是0
    dg = np.maximum(Z, 0)
    dz = np.dot(w2, dz2)*dg
    dw = (np.dot(dz, x.T).T)/m
    db = np.sum(dz, axis=1, keepdims=True)/m
    return dz, dw, db


def Output_layer_Forward_propagation(w, b, X, Y, m):
    # print('A:',X.shape)
    Z2 = np.dot(w.T, X) + b
    A2 = sigmoid(Z2)
    less = (np.sum((Y*np.log(A2)+(1-Y)*np.log(1-A2))))/m
    return A2, less


# A是輸出層預測，也就是輸出,
def Output_layer_Backward_propagation(A2, train_lable, A1, m):
    dz2 = A2-train_lable
    dw2 = np.dot(dz2, A1.T).T/m
    db2 = np.sum(dz2, axis=1, keepdims=True)/m
    # print('dz2:', dz2.shape, 'A2:', A2.shape, 'train_lable:',
    #       train_lable.shape, 'dw2:', dw2.shape, 'A1:', A1.shape, 'db2:', db2.shape, 'm:', m)
    return dz2, dw2, db2


def optimize(w, dw, b, db, learning_rate):
    w -= learning_rate*dw
    b -= learning_rate*db
    return w, b


def logistic_model(train, train_lable, annNum, learning_rate, num_iterations):
    dim = train.shape[0]
    w, b = initWB(dim, annNum)
    w2, b2 = initWB(annNum)
    m = dim
    t0 = time.time()
    for i in range(num_iterations):
        Z, A = Hidden_layer_Forward_propagation(w, b, train)
        A2, less = Output_layer_Forward_propagation(
            w2, b2, A, train_lable, dim)
        dz2, dw2, db2 = Output_layer_Backward_propagation(
            A2, train_lable, A, m)
        dz, dw, db = Hidden_layer_Backward_propagation(train, Z, w2, dz2, m)
        w2, b2 = optimize(w2, dw2, b2, db2, learning_rate)
        w, b = optimize(w, dw, b, db, learning_rate)
        if(i % 100 == 0):
        # if(True):
            t1 = time.time()
            # print('第', i, '次訓練：', '損失函式： ', less, 'w1:', w[1], 'w2:', w2)
            # print('第', i, '次訓練：', '損失函式：', less)
            acc = 0
            total_time = t1-t0
            t0 = time.time()
            for j in range(train_lable.shape[1]):
                yy = train_lable[0][j]
                acc += (A2[0][j]-yy if yy < 1 else yy-A2[0][j])
            acc = (train_lable.shape[1]-acc)*(100/train_lable.shape[1])
            print('第', i, '次訓練：', '損失函式：', less, '準確率:', acc,
                  '%', '時間:', np.round(total_time, 3), 'seconds')
    # 儲存引數
    np.savez('../../myModel.npz', w=w, b=b, w2=w2, b2=b2)


if __name__ == '__main__':
    train, train_lable = init()
    logistic_model(train, train_lable, 20, 0.01, 6000)

    exit()

    imgtest = cv2.resize(cv2.imread('../../test/test/5.jpg'), (weight, height))
    test = imgtest.reshape(1, -1).T/255.
    data = np.load('../../myModel.npz')
    Z = np.dot(data['w'].T, test) + data['b']
    A = relu(Z)
    Z2 = np.dot(data['w2'].T, A) + data['b2']
    A2 = sigmoid(Z2)
    plt.imshow(imgtest)
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.title('本張圖片是貓的概率為%f%%' % (A2*100))
    plt.show()