import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
import utils
import models as ml

%matplotlib inline

"""
Explanation: Have fun with linear model
Author: hdup
My contact info:
hdup huangdan@youhujia.com
evitself evitself@gmail.com
End of explanation
"""


sample_cnt = 100

train_X = np.linspace(-3.0, 3.0, num=sample_cnt, dtype=np.float32).reshape((sample_cnt, 1))
train_y = train_X * 0.375 + 1.1

print(train_X.shape)

"""
Explanation: Let's play with linear regression
Prepare sample data
End of explanation
"""


W, b = ml.create_parameters(1, 1)

# batch learning
for epoch in range(0, 1000):
    h = ml.linear_model(train_X, W, b)
    dW, db = ml.mse_cost_dev(train_X, train_y, h)
    W, b = ml.gd_update(W, b, dW, db, lr=0.01)
    if (epoch + 1) % 100 == 0:
        cur_cost = ml.mse_cost(h, train_y)
        print('epoch: {0}, cost:{1}, W:{2}, b:{3}'.format(epoch + 1, cur_cost, W, b))

# finish
h = ml.linear_model(train_X, W, b)
cost = ml.mse_cost(h, train_y)
print('training finished!')
print('final cost: {0}, W: {1}, b: {2}'.format(cost, W, b))

# then plot some curves
plt.plot(train_X, train_y, 'r+', label='training')
plt.plot(train_X, h, 'b--', label='fitted')
plt.grid(True)
plt.legend()
plt.show()

"""
Explanation: Implement with numpy
1), define model, cost, and derivative

Please see models.py

2), training
End of explanation
"""


train_y_binary = np.array([1.0 if i > 0.5 else 0 for i in train_y]).reshape(sample_cnt, 1)

"""
Explanation: Logistic regression (generalized linear model) with numpy
1), define model, cost, and derivative

Please see models.py

2), prepare label data
End of explanation
"""


W, b = ml.create_parameters(1, 1)

threshold = 0.5

# batch learning
for epoch in range(0, 10000):
    h = ml.logistic_model(train_X, W, b)
    dW, db = ml.log_cost_dev(train_X, train_y_binary, h)
    W, b = ml.gd_update(W, b, dW, db, lr=0.1)
    if (epoch + 1) % 1000 == 0:
        cur_cost = ml.log_cost(h, train_y_binary)
        acc = ml.binary_accuracy(h, train_y_binary, threshold=threshold)
        print('epoch: {0}, cost: {1}, W: {2}, b: {3}, acc: {4}'.format(epoch + 1, cur_cost, W, b, acc))

# finish
h = ml.logistic_model(train_X, W, b)
final_cost = ml.log_cost(h, train_y_binary)
final_acc = ml.binary_accuracy(h, train_y_binary, threshold=threshold)
print('training finished!')
print('final cost: {0}, W: {1}, b: {2}, acc: {3}'.format(final_cost, W, b, final_acc))

# then plot some curves
plt.plot(train_X, train_y_binary, 'r-', label='training')
plt.plot(train_X, h, 'b-', label='fitted')
plt.axhline(y=threshold, color='g', linestyle='-', label='threshold')
plt.grid(True)
plt.ylim([-0.2, 1.2])
plt.legend()
plt.show()

"""
Explanation: 3), training
End of explanation
"""


from helpers import corpus, corpus_len, show_img, prepare_data, one_hot_to_label, showcase

train_data = np.load('./data/data.npz')
train_X=train_data['X']
train_y=train_data['y']
print('train shape X {0}, y {1}'.format(train_X.shape, train_y.shape))

validate_data = np.load('./data/val.npz')
validate_X=validate_data['X']
validate_y=validate_data['y']
print('validation shape X {0}, y {1}'.format(validate_X.shape, validate_y.shape))

test_data = np.load('./data/test.npz')
test_X=test_data['X']
test_y=test_data['y']
print('test shape X {0}, y {1}'.format(test_X.shape, test_y.shape))

train_X, train_y = prepare_data(train_X, train_y)
validate_X, validate_y = prepare_data(validate_X, validate_y)
test_X, test_y = prepare_data(test_X, test_y)

"""
Explanation: Time to have fun with cracking captcha with LR model
1), load and prepare data
End of explanation
"""


sample_cnt = train_X.shape[0]
feature_cnt = train_X.shape[1] * train_X.shape[2]

train_X_a = train_X.reshape((sample_cnt, feature_cnt))
validate_X_a = validate_X.reshape((validate_X.shape[0], feature_cnt))
test_X_a = test_X.reshape((test_X.shape[0], feature_cnt))

train_y_a = train_y[:,0:1]
validate_y_a = validate_y[:,0:1]
test_y_a = test_y[:,0:1]

print(train_X_a.shape)
print(train_y_a.shape)

"""
Explanation: 2), let's take character A as example
You should know one-vs-all
End of explanation
"""


threshold = 0.5

W, b = ml.create_parameters(900, 1)

epoch_arr = []
cost_arr = []
metric_arr = []

step = 20
batch_size = 20

# mini-batch learning
batch_blocks = sample_cnt / batch_size
for epoch in range(0, 100):
    # launch mini-batch
    batch_start = 0    
    while(batch_start + batch_size < sample_cnt):
        batch_X = train_X_a[batch_start:batch_start+batch_size,:]
        batch_y = train_y_a[batch_start:batch_start+batch_size,:] 
        h = ml.logistic_model(batch_X, W, b)
        dW, db = ml.log_cost_dev(batch_X, batch_y, h)
        W, b = ml.gd_update(W, b, dW, db, lr=0.1)
        batch_start += batch_size
        
    # eval epoch
    if (epoch + 1) % step == 0:
        h = ml.logistic_model(validate_X_a, W, b)
        cur_cost = ml.log_cost(h, validate_y_a)
        cur_conf = ml.binary_confusion_matrix(h, validate_y_a, threshold=threshold)        
        print('epoch: {0}, cost: {1}, val_conf: {2}'.format(epoch + 1, cur_cost, cur_conf))
        epoch_arr.append(epoch + 1)
        cost_arr.append(cur_cost)
        metric_arr.append(cur_conf[2])

# finish
predictions = ml.logistic_model(train_X_a, W, b)
final_cost = ml.log_cost(predictions, train_y_a)
final_conf = ml.binary_confusion_matrix(predictions, train_y_a, threshold=threshold)
print('training finished!')
print('final cost: {0}, conf: {1}'.format(final_cost, final_conf))

# calculate test conf
test_h = ml.logistic_model(test_X_a, W, b)
test_cost = ml.log_cost(test_h, test_y_a)
test_conf = ml.binary_confusion_matrix(test_h, test_y_a, threshold=threshold)
print('test cost: {0}, conf: {1}'.format(test_cost, test_conf))

# plot learning curve
fig = plt.figure()

ax1 = fig.add_subplot(111)
ax1.plot(epoch_arr, cost_arr, 'r-', label='cost')
ymax = np.max(cost_arr)
plt.ylim([0.0, ymax * 1.05])
ax1.set_ylabel('cost', color='r')

ax2 = ax1.twinx()
ax2.plot(epoch_arr, metric_arr, 'b-', label='f1')
ymin = np.min(metric_arr)
plt.ylim([ymin-0.02 if ymin-0.02 > 0.0 else 0.0, 1.0])
ax2.set_ylabel('f1', color='b')

plt.axvline(x=0, color='g', linestyle='--')
plt.grid(True)
plt.xlim([-50, epoch + 50])
plt.show()

"""
Explanation: 3), training
End of explanation
"""


sample_count = train_X.shape[0]
feature_count = train_X.shape[1] * train_X.shape[2]
class_count = train_y.shape[1]

train_X = train_X.reshape((sample_count, feature_count))
test_X = test_X.reshape((test_X.shape[0], feature_count))
validate_X = validate_X.reshape((validate_X.shape[0], feature_count))

W, b = ml.create_parameters(feature_count, class_count)

epoch_arr = []
cost_arr = []
metric_arr = []

step = 20
batch_size = 20

# mini-batch learning
batch_blocks = sample_cnt / batch_size
for epoch in range(0, 100):
    # launch mini-batch
    batch_start = 0    
    while(batch_start + batch_size < sample_count):
        batch_X = train_X[batch_start:batch_start+batch_size,:]
        batch_y = train_y[batch_start:batch_start+batch_size,:] 
        h = ml.softmax_regression_model(batch_X, W, b)
        dW, db = ml.crossentropy_cost_dev(batch_X, batch_y, h)
        W, b = ml.gd_update(W, b, dW, db, lr=0.01)
        batch_start += batch_size
        
    # eval epoch
    if (epoch + 1) % step == 0:
        h = ml.softmax_regression_model(validate_X, W, b)
        cur_cost = ml.crossentropy_cost(h, validate_y)  
        cur_acc = ml.categorical_accuracy(h, validate_y)
        print('epoch: {0}, cost: {1}, val_acc: {2}'.format(epoch + 1, cur_cost, cur_acc))
        epoch_arr.append(epoch + 1)
        cost_arr.append(cur_cost)

# finish
h = ml.softmax_regression_model(train_X, W, b)
final_cost = ml.crossentropy_cost(h, train_y)
final_acc = ml.categorical_accuracy(h, train_y)
print('training finished!')
print('final cost: {0}, acc: {1}'.format(final_cost, final_acc))

"""
Explanation: Let's play it with softmax model
End of explanation
"""


