#--------------------------------
#SIMPLE UNIVARIABLE REGRESSION EXAMPLE
#--------------------------------

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from   scipy.optimize import curve_fit

#------------------------
#GENERATE TOY DATA
#------------------------

data_type="E"; #L=linear  E=exp   Q=quadratic
model_type="E" #L=linear  E=exp   Q=quadratic

#GROUND TRUTH DATA GENERATING FUNCTION
def G(x):
	if(data_type=="L"): out=50*x+3
	if(data_type=="Q"): out=100*x**2+50*x+3
	if(data_type=="E"): out=50*np.exp(-0.1*(x-2))+3
	return out

#NOISY DATA
N=200; xmin=-10; xmax=10
x = np.linspace(xmin,xmax,N)
y = G(x)  #PRISTINE DATA
# noise=np.random.normal(loc=0.0, scale=0.05*(max(y)-min(y)),size=len(x))
noise=0.15*(max(y)-min(y))*np.random.uniform(-1,1,size=len(x))
yn = y+ noise

#GROUND TRUTH DATA
xe = np.linspace(xmin,2*xmax,int(2*N))
ye = G(xe)

#NORMALIZE
xe=(xe-np.mean(x))/np.std(x)
ye=(ye-np.mean(y))/np.std(y)
yn=(yn-np.mean(y))/np.std(y)

x=(x-np.mean(x))/np.std(x)
y=(y-np.mean(y))/np.std(y)


#------------------------
#TRAINING
#------------------------

##FITTING MODEL 
def model(x,p0,p1,p2,p3):
	if(model_type=="L"): out=p1*x+p0
	if(model_type=="Q"): out=p2*x**2+p1*x+p0
	if(model_type=="E"): out=p3*np.exp(p2*(x-p1))+p0
	return out

#INITIAL GUESS FOR PARAMETERS p0,p1,p2,p3
po=[0.1,0.1,0.1,0.1]

#FIT MODEL
popt, pcov = curve_fit(model, x, yn,po) #,[0.1,0.1,0.1])

#------------------------
#PLOT
#------------------------

fig, ax = plt.subplots()
ax.plot(x, yn, 'o', label='Data')
ax.plot(xe, ye, '-', label='Ground-Truth')
ax.plot(xe, model(xe, *popt), 'r-', label="Model")

ax.legend()
FS=18   #FONT SIZE
plt.xlabel('X', fontsize=FS)
plt.ylabel('Y', fontsize=FS)

plt.show()

