# Import some libraries that will be necessary for working with data and displaying plots

# To visualize plots in the notebook
%matplotlib inline 

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy.io       # To read matlab files
import pylab

"""
Explanation: Parametric Model-Based regression
Notebook version: 1.3 (Sep 20, 2019)

Author: Jesús Cid-Sueiro (jesus.cid@uc3m.es)
        Jerónimo Arenas García (jarenas@tsc.uc3m.es)

Changes: v.1.0 - First version, expanding some cells from the Bayesian Regression 
                 notebook
         v.1.1 - Python 3 version.
         v.1.2 - Revised presentation. 
         v.1.3 - Updated index notation

Pending changes: * Include regression on the stock data
End of explanation
"""


# wML = <FILL IN>

print("The ML estimate is {}".format(wML))

"""
Explanation: A quick note on the mathematical notation
In this notebook we will make extensive use of probability distributions. In general, we will use capital leters
${\bf X}$, $S$, $E$ ..., to denote random variables, and lower-case letters ${\bf x}$, $s$, $\epsilon$ ..., to denote the values they can take. 
In general, we will use letter $p$ for probability density functions (pdf). When necessary, we will use, capital subindices to make the random variable explicit. For instance, $p_{{\bf X}, S}({\bf x}, s)$ would be the joint pdf of random variables ${\bf X}$ and $S$ at values ${\bf x}$ and $s$, respectively. 
However, to avoid a notation overload, we will omit subindices when they are clear from the context. For instance, we will use $p({\bf x}, s)$ instead of $p_{{\bf X}, S}({\bf x}, s)$.
1. Model-based parametric regression
1.1. The regression problem
Given an observation vector ${\bf x}$, the goal of the regression problem is to find a function $f({\bf x})$ providing good predictions about some unknown variable $s$. To do so, we assume that a set of labelled training examples, ${{\bf x}k, s_k}{k=0}^{K-1}$ is available. 
The predictor function should make good predictions for new observations ${\bf x}$ not used during training. In practice, this is tested using a second set (the test set) of labelled samples.
1.2. The underlying model assumption
Many regression algorithms are grounded on the idea that all samples from the training set have been generated 
independently by some common stochastic process.
<img src="figs/data_model.png" width=180>
If $p({\bf x}, s)$ were known, we could apply estimation theory to estimate $s$ for a given ${\bf x}$ using $p$. For instance, we could apply any of the following classical estimates:

Maximum A Posterior (MAP): $$\hat{s}_{\text{MAP}} = \arg\max_s p(s| {\bf x})$$
Minimum Mean Square Error (MSE): $$\hat{s}_{\text{MSE}} = \mathbb{E}{S |{\bf x}} = \int s \, p(s| {\bf x}) \, ds $$

Note that, since these estimators depend on $p(s |{\bf x})$, knowing the posterior distribution of the target variable is enough, and we do not need to know the joint distribution $p({\bf x}, s)$.
More importantly, note that if we knew the underlying model, we would not need the data in ${\cal D}$ to make predictions on new data.
Exercise 1:
Assume the target variable $s$ is a scaled noisy version of the input variable $x$: 
$$
s = 2 x + \epsilon
$$
where $\epsilon$ is Gaussian a noise variable with zero mean and unit variance, which does not depend on $x$.

Compute the target model $p(s| x)$
Compute prediction $\hat{s}_\text{MAP}$ for an arbitrary input $x$
Compute prediction $\hat{s}_\text{MSE}$ for an arbitrary input $x$
Compute prediction $\hat{s}_\text{MSE}$ for input $x=4$

Solution:
1.3. Model-based regression
In practice, the underlying model is usually unknown. 
Model based-regression methods exploit the idea of using the training data to estimate the posterior distribution $p(s|{\bf x})$ and then apply estimation theory to make predictions.
<img src="figs/ModelBasedReg.png" width=280>
1.4. Parametric model-based regression
In some cases, we may have a partial knowledge about the underlying mode. In this notebook we will assume that $p$ belongs to a parametric family of distributions $p(s|{\bf x},{\bf w})$, where ${\bf w}$ is some unknown parameter. 
Exercise 2:
Assume the target variable $s$ is a scaled noisy version of the input variable $x$: 
$$
s = w x + \epsilon
$$
where $\epsilon$ is Gaussian a noise variable with zero mean and unit variance, which does not depend on $x$. Assume that $w$ is known. 
  1. Compute the target model $p(s| x, w)$
  2. Compute prediction $\hat{s}\text{MAP}$ for an arbitrary input $x$
  3. Compute prediction $\hat{s}\text{MSE}$ for an arbitrary input $x$
Solution:
We will use the training data to estimate ${\bf w}$
<img src="figs/ParametricReg.png" width=300>
The estimation of ${\bf w}$ from a given dataset $\mathcal{D}$ is the goal of the following sections
2. Maximum Likelihood parameter estimation.
The ML (Maximum Likelihood) principle is well-known in statistics and can be stated as follows: take the value of the parameter to be estimated (in our case, ${\bf w}$) that best explains the given observations (in our case, the training dataset $\mathcal{D}$). Mathematically, this can be expressed as follows:
$$
\hat{\bf w}{\text{ML}} = \arg \max{\bf w} p(\mathcal{D}|{\bf w})
$$
Exercise 3:
All samples in dataset ${\cal D} = {(x_k, s_k), k=0,\ldots,K-1 }$ 
$$
s_k = w \cdot x_k + \epsilon_k
$$
where $\epsilon_k$ are i.i.d. (independent and identically distributed) Gaussian noise random variables with zero mean and unit variance, which do not depend on $x_k$. 
Compute the ML estimate, $\hat{w}_{\text{ML}}$, of $w$.
Solution:
</SOL>

4.2. Compute the ML estimate
End of explanation
"""


sigma_eps = 1
K = len(s)
wGrid = np.arange(-0.5, 2, 0.01)

p = []
for w in wGrid:
    d = s - X*w
    # p.append(<FILL IN>)

# Compute the likelihood for the ML parameter wML
# d = <FILL IN>
# pML = [<FILL IN>]

# Plot the likelihood function and the optimal value
plt.figure()
plt.plot(wGrid, p)
plt.stem([wML], pML)
plt.xlabel('$w$')
plt.ylabel('Likelihood function')
plt.show()

"""
Explanation: 4.3. Plot the likelihood as a function of parameter $w$ along the interval $-0.5\le w \le 2$, verifying that the ML estimate takes the maximum value.
End of explanation
"""


xgrid = np.arange(0, 1.2, 0.01)
# sML = <FILL IN>

plt.figure()
plt.scatter(X, s)
# plt.plot(<FILL IN>)
plt.xlabel('x')
plt.ylabel('s')
plt.axis('tight')
plt.show()

"""
Explanation: 4.4. Plot the prediction function on top of the data scatter plot
End of explanation
"""


X = np.array([0.15, 0.41, 0.53, 0.80, 0.89, 0.92, 0.95]) 
s = np.array([0.09, 0.16, 0.63, 0.44, 0.55, 0.82, 0.95]) 

"""
Explanation: 2.1. Model assumptions
In order to solve exercise 4 we have taken advantage of the statistical independence of the noise components. Some independence assumptions are required in general to compute the ML estimate in other scenarios.
In order to estimate ${\bf w}$ from the training data in a mathematicaly rigorous and compact form let us group the target variables into a vector
$$
{\bf s} = \left(s_0, \dots, s_{K-1}\right)^\top
$$
and the input vectors into a matrix
$$
{\bf X} = \left({\bf x}0, \dots, {\bf x}{K-1}\right)^\top
$$
We will make the following assumptions:

A1. All samples in ${\cal D}$ have been generated by the same distribution, $p({\bf x}, s \mid {\bf w})$
A2. Input variables ${\bf x}$ do not depend on ${\bf w}$. This implies that
$$
p({\bf X} \mid {\bf w}) = p({\bf X})
$$
A3. Targets $s_{0},\ldots, s_{K-1}$ are statistically independent, given ${\bf w}$ and the inputs ${\bf x}0,\ldots, {\bf x}{K-1}$, that is:
$$
p({\bf s} \mid {\bf X}, {\bf w}) = \prod_{k=0}^{K-1} p(s_k \mid {\bf x}_k, {\bf w})
$$

Since ${\cal D} = ({\bf X}, {\bf s})$, we can write
$$p(\mathcal{D}|{\bf w}) 
    = p({\bf s}, {\bf X}|{\bf w}) 
    = p({\bf s} | {\bf X}, {\bf w}) p({\bf X}|{\bf w})
$$
Using assumption A2,
$$
p(\mathcal{D}|{\bf w}) 
    = p({\bf s} | {\bf X}, {\bf w}) p({\bf X})
$$
and, finally, using assumption A3, we can express the estimation problem as the computation of
\begin{align}
\hat{\bf w}{\text{ML}}
                &= \arg \max{\bf w} p({\bf s}|{\bf X},{\bf w})  \
   \qquad \quad &= \arg \max_{\bf w} \prod_{k=0}^{K-1} p(s_k \mid {\bf x}k, {\bf w})  \
   \qquad \quad &= \arg \max{\bf w} \sum_{k=0}^{K-1}\log  p(s_k \mid {\bf x}_k, {\bf w})
\end{align}
Any of the last three terms can be used to optimize ${\bf w}$. The sum in the last term is usually called the log-likelihood function, $L({\bf w})$, whereas the product in the previous line is simply referred as the likelihood function.
2.2. Summary.
Let's summarize what we need to do in order to design a regression algorithm based on ML estimation:

Assume a parametric data model $p(s| {\bf x},{\bf w})$
Using the data model and the i.i.d. assumption, compute $p({\bf s}| {\bf X},{\bf w})$.
Find an expression for ${\bf w}_{\text{ML}}$
Assuming ${\bf w} = {\bf w}_{\text{ML}}$, compute the MAP or the minimum MSE estimate of $s$ given ${\bf x}$.

3. ML estimation for a Gaussian model.
3.1. Step 1: The Gaussian generative model
Let us assume that the target variables $s_k$ in dataset $\mathcal{D}$ are given by
$$
s_k = {\bf w}^\top {\bf z}_k + \varepsilon_k
$$
where ${\bf z}k$ is the result of some transformation of the inputs, ${\bf z}_k = T({\bf x}_k)$, and $\varepsilon_k$ are i.i.d. instances of a Gaussian random variable with mean zero and varianze $\sigma\varepsilon^2$, i.e.,
$$
p_E(\varepsilon) = \frac{1}{\sqrt{2\pi}\sigma_\varepsilon}
                   \exp\left(-\frac{\varepsilon^2}{2\sigma_\varepsilon^2}\right)
$$
Assuming that the noise variables are independent of ${\bf x}$ and ${\bf w}$, then, for a given ${\bf x}$ and ${\bf w}$, the target variable is gaussian with mean ${\bf w}^\top {\bf z}k$ and variance $\sigma\varepsilon^2$
$$
p(s|{\bf x}, {\bf w}) = p_E(s-{\bf w}^\top{\bf z}) =
    \frac{1}{\sqrt{2\pi}\sigma_\varepsilon}
    \exp\left(-\frac{(s-{\bf w}^\top{\bf z})^2}{2\sigma_\varepsilon^2}\right)
$$
3.2. Step 2: Likelihood function
Now we need to compute the likelihood function $p({\bf s}, {\bf X} | {\bf w})$. If the samples are i.i.d. we can write
$$
p({\bf s}| {\bf X}, {\bf w})
    = \prod_{k=0}^{K-1} p(s_k| {\bf x}k, {\bf w}) 
    = \prod{k=0}^{K-1} \frac{1}{\sqrt{2\pi}\sigma_\varepsilon}
    \exp\left(-\frac{\left(s_k-{\bf w}^\top{\bf z}k\right)^2}{2\sigma\varepsilon^2}\right) \
    = \left(\frac{1}{\sqrt{2\pi}\sigma_\varepsilon}\right)^K
      \exp\left(-\sum_{k=1}^K \frac{\left(s_k-{\bf w}^\top{\bf z}k\right)^2}{2\sigma\varepsilon^2}\right) \
$$
Finally, grouping variables ${\bf z}k$ in
$${\bf Z} = \left({\bf z}_0, \dots, {\bf z}{K-1}\right)^\top$$
we get
$$
p({\bf s}| {\bf X}, {\bf w})
    = \left(\frac{1}{\sqrt{2\pi}\sigma_\varepsilon}\right)^K
      \exp\left(-\frac{1}{2\sigma_\varepsilon^2}\|{\bf s}-{\bf Z}{\bf w}\|^2\right)
$$
3.3. Step 3: ML estimation.
The <b>maximum likelihood</b> solution is then given by:
$$
{\bf w}{ML} = \arg \max{\bf w} p({\bf s}|{\bf w}) = \arg \min_{\bf w} \|{\bf s} - {\bf Z}{\bf w}\|^2
$$
Note that $\|{\bf s} - {\bf Z}{\bf w}\|^2$ is the sum or the squared prediction errors (Sum of Squared Errors, SSE) for all samples in the dataset. This is also called the * Least Squares * (LS) solution.
The LS solution can be easily computed by differentiation,
$$
\nabla_{\bf w} \|{\bf s} - {\bf Z}{\bf w}\|^2\Bigg|{{\bf w} = {\bf w}\text{ML}}
= - 2 {\bf Z}^\top{\bf s} + 2 {\bf Z}^\top{\bf Z} {\bf w}{\text{ML}}
= {\bf 0}
$$
and it is equal to
$$
{\bf w}\text{ML} = ({\bf Z}^\top{\bf Z})^{-1}{\bf Z}^\top{\bf s}
$$
3.4. Step 4: Prediction function.
The last step consists on computing an estimate of $s$ by assuming that the true value of the weight parameters is ${\bf w}\text{ML}$. In particular, the minimum MSE estimate is
$$
\hat{s}\text{MSE} = \mathbb{E}{s|{\bf x},{\bf w}_\text{ML}}
$$
Knowing that, given ${\bf x}$ and ${\bf w}$, $s$ is normally distributed with mean ${\bf w}^\top {\bf z}$ we can write
$$
\hat{s}\text{MSE} = {\bf w}\text{ML}^\top {\bf z}
$$
Exercise 5:
Assume that the targets in the one-dimensional dataset given by
End of explanation
"""


sigma_eps = 0.3

"""
Explanation: have been generated by the polynomial Gaussian model 
$$
s = w_0 + w_1 x + w_2 x^2 + \epsilon
$$
(i.e., with ${\bf z} = T(x) = (1, x, x^2)^\intercal$) with noise variance
End of explanation
"""


# Compute the extended input matrix Z
nx = len(X)
# Z = <FILL IN>

# Compute the ML estimate using linalg.lstsq from Numpy.
# wML = <FILL IN>
print(wML)


"""
Explanation: 5.1. Compute the ML estimate.
End of explanation
"""


K = len(s)

# Compute the likelihood for the ML parameter wML
# d = <FILL IN>

# LwML = [<FILL IN>]

print(LwML)

"""
Explanation: 5.2. Compute the value of the log-likelihood function for ${\bf w}={\bf w}_\text{ML}$.
End of explanation
"""


xgrid = np.arange(0, 1.2, 0.01)
nx = len(xgrid)
# Compute the input matrix for the grid data in x
# Z = <FILL IN>

# sML = <FILL IN>

plt.figure()
plt.scatter(X, s)
# plt.plot(<FILL IN>)
plt.xlabel('x')
plt.ylabel('s')
plt.axis('tight')
plt.show()

"""
Explanation: 5.3. Plot the prediction function over the data scatter plot
End of explanation
"""


K = len(s)
wGrid = np.arange(0, 6, 0.01)

p = []
Px = np.prod(X)
xs = np.dot(X,s)
for w in wGrid:
    # p.append(<FILL IN>)

plt.figure()
# plt.plot(<FILL IN>)
plt.xlabel('$w$')
plt.ylabel('Likelihood function')
plt.show()

"""
Explanation: Exercise 6:
Assume the dataset $\mathcal{D} = {(x_k, s_k, k=0,\ldots, K-1}$ contains i.i.d. samples from a distribution with posterior density given by
$$
p(s \mid x, w) = w x \exp(- w x s), \qquad s\ge0, \,\, x\ge 0, \,\, w\ge 0
$$

6.1. Determine an expression for the likelihood function

Solution:
<SOL>
</SOL>

6.2. Draw the likelihood function for the dataset in Exercise 4 in the range $0\le w\le 6$.
End of explanation
"""


# wML = <FILL IN>
print(wML)

"""
Explanation: 6.3. Determine the maximum likelihood coefficient, $w_\text{ML}$. 

(Hint: you can maximize the log-likelihood function instead of the likelihood function in order to simplify the differentiation)
Solution:
<SOL>
</SOL>

6.4. Compute $w_\text{ML}$ for the dataset in Exercise 4
End of explanation
"""


xgrid = np.arange(0.1, 1.2, 0.01)
# sML = <FILL IN>

plt.figure()
plt.scatter(X, s)
# plt.plot(<FILL IN>)
plt.xlabel('x')
plt.ylabel('s')
plt.axis('tight')
plt.show()

"""
Explanation: 6.5. Assuming $w = w_\text{ML}$, compute the prediction function based on the estimate $s_{MSE}$

Solution:
<SOL>
</SOL>

6.6. Plot the prediction function obtained in ap. 6.5, and compare it with the linear predictor in exercise 4
End of explanation
"""


