import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt

import tensorflow as tf
import Henbun as hb

"""
Explanation: Structure of Henbun
This notebook describes structure of Henbun and how to construct a model with it.
We construct a linear model as an example.
Note that Henbun is not very effective for such a very simple model.
This example is just to present Henbun's structure and its main functionalities.
Import libraries
End of explanation
"""


X = np.linspace(0,6,40).reshape(-1,1)
Y = 0.5*X + np.random.randn(40,1)*0.3 + 0.4

plt.figure(figsize=(5,3))
plt.plot(X, Y, 'o', label='data')
plt.plot(X, 0.5*X+0.4, '--k', label='true')
plt.legend(loc='best')

"""
Explanation: Toy data
End of explanation
"""


class LinearFit(hb.model.Model):
    def setUp(self):
        """
        This method is called just after the instanciation of this class.
        In this method, we need to define any variables and data to be used in this model.
        """
        # data should be stored in hb.param.Data class
        self.X = hb.param.Data(X)
        self.Y = hb.param.Data(Y)
        
        # Parameters that are defined as hb.param.Variable will be optimized by this model.
        self.a = hb.param.Variable(shape=[1])
        self.b = hb.param.Variable(shape=[1])
        
    @hb.model.AutoOptimize()
    def MinusLoss(self):
        """
        Any method decorated by @AutoOptimize can be optimized.
        Here we return the minus of Loss
        """
        return -tf.reduce_sum(tf.square(self.Y - (self.a + self.b * self.X)))

linear_fit = LinearFit()

"""
Explanation: Linear Fit Model
We assume the data $y_i$ has a linear relation with input $x_i$.
$$y_i \sim a + b x_i$$
with coefficients a and b.
We seek $a$ and $b$ that minimizes the following loss
$$
\mathrm{Loss} = \sum_i{(y_i - (a + b * x_i))^2}
$$
Construct a linear fit model
End of explanation
"""


# To initialize values call .initialize()
linear_fit.initialize()

print(linear_fit.a.value, linear_fit.b.value)

# Manually initialize variables
# To initialize manually, just set the desired value,
linear_fit.a = 0.1
linear_fit.b = 0.1
# To reflect this operation, another `initialize` call is necessary
linear_fit.initialize()

# see the values were updated.
print(linear_fit.a.value, linear_fit.b.value)

"""
Explanation: See what is inside in Variables
The Variables are automatically initialized.
To see what values are inside this object, .value property can be used after the initialization.
End of explanation
"""


# compilation
linear_fit.MinusLoss().compile()

# optimization can be made after .compile() method is completed.
linear_fit.MinusLoss().optimize(maxiter=10000)

"""
Explanation: Find parameters $a$ and $b$ that minimizes the Loss
Compilation and optimization
To make an optimization, we need to compile the objective.
In our case, the objective is linear_fit.MinusLoss.
End of explanation
"""


print(linear_fit.a.value, linear_fit.b.value)

plt.figure(figsize=(5,3))
plt.plot(X, Y, 'o', label='data')
plt.plot(X, 0.5*X+0.4, '--k', label='true')
plt.plot(X, linear_fit.a.value+linear_fit.b.value*X, '-r', lw=2, alpha=0.5, label='fit')
plt.plot()
plt.legend(loc='best')

"""
Explanation: See the results
The values in Variable object is updated by optimize method.
End of explanation
"""


# Construct a probabilistic linear model
class LinearModel(hb.model.Model):
    def setUp(self):
        # data should be stored in hb.param.Data class
        self.X = hb.param.Data(X)
        self.Y = hb.param.Data(Y)
        
        # Parameters that are defined as hb.param.Variable will be optimized by this model.
        self.a = hb.param.Variable(shape=[1])
        self.b = hb.param.Variable(shape=[1])
        
        # Addition to linear_model, we define the variance parameter.
        # This parameter should be positive. 
        # It can be achieved by passing transform option.
        self.sigma = hb.param.Variable(shape=[1], transform=hb.transforms.positive)
        
    @hb.model.AutoOptimize()
    def logp(self):
        """
        This method returns the sum of log-likelihood and log-prior.
        """
        log_lik = hb.densities.gaussian(self.Y, self.a + self.b * self.X, self.sigma)
        log_prior = hb.densities.gaussian(self.a, 0.0, 1.0)\
                  + hb.densities.gaussian(self.b, 0.0, 1.0)\
                  + hb.densities.gaussian(self.sigma, 0.0, 1.0)
        
        return tf.reduce_sum(log_lik) + log_prior

"""
Explanation: Probabilistic linaer model
Here, we construct a similar linear model within Bayesian framework.
The likelihood is assumed as Gaussian,
$$
p(y_i | a, b, x_i) = \mathcal{N}(y_i | a + b*x_i, \sigma)
$$
where $\sigma$ is the variance parameter.
We assume weak prior for $a$, $b$ and $\sigma$,
$$
p(a) = \mathcal{N}(a|0,1) \
p(b) = \mathcal{N}(b|0,1) \
p(\sigma) = \mathcal{N}(\sigma|0,1) \
$$
In this model, we will find $a$ and $b$ that jointly maximizes the posterior distribution,
$$
p(a, b, \sigma| \mathbf{x}, \mathbf{y}) 
\propto 
\prod_{i}\mathcal{N}(y_i | a + b*x_i, \sigma) 
\mathcal{N}(a|0,1) \mathcal{N}(b|0,1)  \mathcal{N}(\sigma|0,1)
$$
End of explanation
"""


# similary, we compile and optimize optimize the model.
plinear_model = LinearModel()
plinear_model.logp().compile()
plinear_model.logp().optimize(maxiter=10000)

print(plinear_model.a.value, plinear_model.b.value, plinear_model.sigma.value)

plt.figure(figsize=(5,3))
plt.plot(X, Y, 'o', label='data')
plt.plot(X, 0.5*X+0.4, '--k', label='true')
plt.plot(X, plinear_model.a.value+plinear_model.b.value*X, '-r', lw=2, alpha=0.5, label='fit')
plt.plot()
plt.legend(loc='best')

"""
Explanation: Result
End of explanation
"""


type(linear_fit.a._tensor)

"""
Explanation: What is inside hb.param.Variable
hb.param.Variable is an object that wraps tf.Variable.
In fact, hb.param.Variable._tensor is tf.Variable
End of explanation
"""


print(plinear_model.sigma._tensor) # <- parameters that spans in real space
print(plinear_model.sigma.tensor()) # <- transformed parameters that is cast to the positive space

"""
Explanation: Variable object has .tensor() method, that returns the transformed tensor.
End of explanation
"""


print('Without tf_mode: '+str(type(plinear_model.a)))
with plinear_model.tf_mode():
    print('With tf_mode: ' + str(type(plinear_model.a)))

"""
Explanation: Any parameterized object such as hb.model.Model have .tf_mode().
Within tf_mode, Variable object is seen as its .tensor() method.
End of explanation
"""


print(type(plinear_model.logp()))

optimizer = plinear_model.logp()
optimizer.__dict__

"""
Explanation: In decorated methods by @hb.model.AutoOptimize(), 
tf_mode is automatically switched on.
What is inside @hb.model.AutoOptimize()
The methods decorated by @hb.model.AutoOptimize() returns 
hb.model.Optimizer object, 
that contains the objective function, variables to be minimized, and so on.
End of explanation
"""


