# Imports
import numpy as np
import matplotlib.pyplot as plt
import scipy as sp
import scipy.interpolate
import pandas as pd

# scikit
import sklearn
import sklearn.ensemble
import sklearn.metrics
import sklearn.utils
import sklearn.cross_validation
import sklearn.isotonic

plt.ion()

# Simulate some true and observed data
offset = 33
slope1 = 10
slope2 = 2
points = 1000
x = np.linspace(0, 100, points)
y_true = 100 - ((x < offset) * x * slope1 + (x >= offset) * ((x-offset) * slope2 + (offset * slope1)))
y_observed = y_true + (0.5 - np.random.random(y_true.shape)) * np.sin(100-x) * (100-x) * 2

%matplotlib inline
plt.figure(figsize=(10, 8))
plt.title('Isotonic Example')
plt.xlabel('x')
plt.ylabel('y')
plt.scatter(x, y_true, alpha=0.5, color='blue')
plt.scatter(x, y_observed, alpha=0.5, color='red')
plt.legend(('True', 'Observed'))

# Fit isotonic
ir = sklearn.isotonic.IsotonicRegression(increasing='auto')
ir.fit(x, y_observed)
y_ir = ir.predict(x)

# Fit a spline with FITPACK
fitpack_spline = sp.interpolate.UnivariateSpline(x, y_observed)
y_fitpack = fitpack_spline(x)

%matplotlib inline
plt.figure(figsize=(10, 8))
plt.title('Isotonic Example')
plt.xlabel('x')
plt.ylabel('y')
plt.scatter(x, y_true, alpha=0.5, color='blue')
plt.scatter(x, y_observed, alpha=0.1, color='red')
plt.scatter(x, y_fitpack, alpha=0.25, color='purple')
plt.scatter(x, y_ir, alpha=0.5, color='green')
plt.legend(('True', 'Observed', 'FITPACK', 'Isotonic'))

# Plot the final function z

# Sample new X and Y data
x_new = np.random.randint(0, 100, points)
y_true_new = 100 - ((x_new < offset) * x_new * slope1 + (x_new >= offset) * ((x_new-offset) * slope2 + (offset * slope1)))
y_observed_new = y_true_new + (0.5 - np.random.random(y_true_new.shape)) * np.sin(100-x_new) * (100-x_new) * 2
y_ir_new = ir.predict(x_new)
y_fitpack_new = fitpack_spline(x_new)

# Plot resulting true and estimated
plt.figure(figsize=(10,8))
plt.title('Function with isotonic smoothing')
plt.xlabel('y')
plt.ylabel('z')
plt.scatter(x_new, -2 * y_true_new**3 + y_true_new**2 + y_true_new * 3.5 + 9, alpha=0.5, color='blue')
plt.scatter(x_new, -2 * y_observed_new**3 + y_observed_new**2 + y_observed_new * 3.5 + 9, alpha=0.1, color='red')
plt.scatter(x_new, -2 * y_fitpack_new**3 + y_fitpack_new**2 + y_fitpack_new * 3.5 + 9, alpha=0.2, color='purple')
plt.scatter(x_new, -2 * y_ir_new**3 + y_ir_new**2 + y_ir_new * 3.5 + 9, alpha=0.5, color='green')
plt.legend(('True', 'Observed', 'FITPACK', 'Isotonic'))