#!/usr/bin/env python
# -*- coding:utf-8 -*-

import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets, linear_model
import theano
import theano.tensor as T
import pdb

# Load the diabetes dataset
diabetes = datasets.load_diabetes()

# Use only one feature
diabetes_X = diabetes.data[:, np.newaxis, 2]

# Split the data into training/testing sets
diabetes_X_train = diabetes_X[:-20]
diabetes_X_test = diabetes_X[-20:]

# Split the targets into training/testing sets
diabetes_y_train = diabetes.target[:-20]
diabetes_y_test = diabetes.target[-20:]
diabetes_X_test = diabetes_X_test.reshape(-1, 1)
diabetes_y_test = diabetes_y_test.reshape(-1, 1)
diabetes_X_train = diabetes_X_train.reshape(-1, 1)
diabetes_y_train = diabetes_y_train.reshape(-1, 1)

x = T.dmatrix('x')
y = T.dmatrix('y')

W = theano.shared(value=np.zeros((1,1), dtype=theano.config.floatX), 
                    name='W', borrow=True)
b = theano.shared(value=np.zeros((1,), dtype=theano.config.floatX)+0.1, 
                    name='b', borrow=True)

p_y_given_x = T.dot(x,W) + b

cost = T.mean(T.sqr(p_y_given_x - y))

g_W = T.grad(cost=cost, wrt=W)
g_b = T.grad(cost=cost, wrt=b)

learning_rage = 0.1
updates = [(W, W - learning_rage * g_W),
           (b, b - learning_rage * g_b)]
# pdb.set_trace()
train_model = theano.function(inputs=[x, y], outputs=cost, updates=updates)

predict = theano.function(inputs=[x], outputs=p_y_given_x)

test_err = theano.function(inputs=[x, y], outputs=cost)

for i in range(10000):
    # training
    err = train_model(diabetes_X_train, diabetes_y_train)
    if i % 500 == 0:
        print(err)

# Create linear regression object
regr = linear_model.LinearRegression()

# Train the model using the training sets
regr.fit(diabetes_X_train, diabetes_y_train)

# The coefficients
print('Sklearn Coefficients: \n', regr.coef_)
print('Sklearn intercept_: \n', regr.intercept_)
print('Theano W:', W.get_value())
print('Theano b:', b.get_value())
# The mean squared error
print("Sklearn Mean squared error: %.2f"
      % np.mean((regr.predict(diabetes_X_test) - diabetes_y_test) ** 2))

print("Theano Mean squared error: %.2f" 
      % test_err(diabetes_X_test, diabetes_y_test))

# Explained variance score: 1 is perfect prediction
print('Variance score: %.2f' % regr.score(diabetes_X_test, diabetes_y_test))

# Plot outputs
plt.scatter(diabetes_X_test, diabetes_y_test,  color='black')
plt.plot(diabetes_X_test, predict(diabetes_X_test), color='blue',
         linewidth=3)
plt.plot(diabetes_X_test, regr.predict(diabetes_X_test), color='red',
         linewidth=3)

plt.xticks(())
plt.yticks(())

plt.show()