# -*- coding: utf-8 -*-
"""
# @Time    : 2019/11/10 11:16
# @Author  : zhaobinghao
# @File    : assignment4.py
"""
from sklearn.model_selection import train_test_split
from sklearn import metrics
from sklearn import datasets
import matplotlib.pyplot as plt
import numpy as np


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


def initialize_parameters(dim):
    # w = np.random.rand(dim, 1)
    w = np.random.rand(dim, 1)
    b = 0
    assert (w.shape == (dim, 1))
    assert (isinstance(b, float) or isinstance(b, int))
    return w, b


def propagate(w, b, X, Y):
    '''
    Implement the cost function and its gradient for the propagation

    Arguments:
    w - weights
    b - bias
    X - data
    Y - ground truth
    '''
    Y = np.array(Y).reshape(-1, 1)
    m = X.shape[1]
    A = sigmoid(np.dot(X, w) + b)
    cost = -np.mean(np.nan_to_num(Y * np.log(A) + (1 - Y) * np.log(1 - A)))

    dw = np.dot(X.T, (A - Y)) / m
    db = np.sum(A - Y) / m

    assert (dw.shape == w.shape)
    assert (db.dtype == float)
    cost = np.squeeze(cost)
    assert (cost.shape == ())

    grads = {'dw': dw,
             'db': db}
    return grads, cost


def optimize(w, b, X, Y, num_iterations, learning_rate, print_cost=False):
    '''
    This function optimize w and b by running a gradient descen algorithm

    Arguments:
    w - weights
    b - bias
    X - data
    Y - ground truth
    num_iterations -- number of iterations of the optimization loop
    learning_rate -- learning rate of the gradient descent update rule
    print_cost -- True to print the loss every 100 steps

    Returns:
    params - dictionary containing the weights w and bias b
    grads -- dictionary containing the gradients of the weights and bias with respect to the cost function
    costs -- list of all the costs computed during the optimization, this will be used to plot the learning curve.

    '''

    costs = []

    for i in range(num_iterations):

        grads, cost = propagate(w, b, X, Y)

        dw = grads['dw']
        db = grads['db']

        w = -learning_rate * dw + w
        b = -learning_rate * db + b

        if i % 100 == 0:
            costs.append(cost)
        if print_cost and i % 100 == 0:
            print("Cost after iteration %i: %f" % (i, cost))

    params = {"w": w,
              "b": b}

    grads = {"dw": dw,
             "db": db}

    return params, grads, costs


def predict(w, b, X):
    '''
    Predict whether the label is 0 or 1 using learned logistic regression parameters (w, b)

    Arguments:
    w -- weights
    b -- bias
    X -- data

    Returns:
    Y_prediction -- a numpy array (vector) containing all predictions (0/1) for the examples in X
    '''
    m = X.shape[1]
    w = w.reshape(m, 1)
    Y_prediction = sigmoid(np.dot(X, w) + b)
    assert (Y_prediction.shape == (X.shape[0], 1))
    return Y_prediction


def result_report(y_true, y_pred, thresh_hold=0.5):
    y = [1 if y > thresh_hold else 0 for y in y_pred]
    print('confusion matrix')
    print(metrics.confusion_matrix(y_true, y))
    print('accuracy score', metrics.accuracy_score(y_true, y))
    print('precision score: ', metrics.precision_score(y_true, y))
    print('recall score: ', metrics.recall_score(y_true, y))
    print('f1 score: ', metrics.f1_score(y_true, y))


def model(X_train, y_train, X_test, y_test, num_iterations, learning_rate, print_cost):
    w, b = initialize_parameters(X_train.shape[1])
    params, grads, costs = optimize(w, b, X_train, y_train, num_iterations, learning_rate, print_cost=True)
    # plt.plot(range(len(costs)), costs)
    # plt.show()
    # train
    print('train', '_' * 50)
    y_pred_train = predict(params['w'], params['b'], X_train)
    result_report(y_train, y_pred_train)
    # test
    print('train', '_' * 50)
    y_pred_test = predict(params['w'], params['b'], X_test)
    result_report(y_test, y_pred_test)
    return {'w': params['w'],
            'b': params['b'],
            'cost': costs}


if __name__ == '__main__':
    # Loading the data
    digits = datasets.load_digits()
    # Split the data into training set and test set
    data = digits.data / np.max(digits.data)  # 归一化
    X_train, X_test, y_train, y_test = train_test_split(data, digits.target, test_size=0.25)
    y_train[y_train < 5] = 0
    y_train[y_train >= 5] = 1
    y_test[y_test < 5] = 0
    y_test[y_test >= 5] = 1

    result = model(X_train, y_train, X_test, y_test, num_iterations=2000, learning_rate=0.05, print_cost=True)
