'''
Simplified version of cointoss.py to simply do MLE of the heads-probability.

Coin tosses :math:`x` are generated by a Bernoulli process with unknown heads probability :math:`z = P(x=1)`,
so z is the latent/hidden variable, and we want to find the Max-(log)-Likelihood estimate of z,
based on observations of the tosses x.
'''
import numpy as np
import scipy.stats as stats
from scipy.special import gamma
import torch as t
import torch.nn as nn
from torch.autograd import Variable, backward
import torch.optim as optim
import seaborn as sns
import matplotlib.pyplot as plt

def gen_tosses(z, n=1):
    return Variable(t.FloatTensor(np.random.binomial(1, z, n)*1.0).unsqueeze(1))

def logit(x):
    return t.log(x/(1-x))

def sigmoid(x):
    return 1.0/(1 + np.exp(-x))


# MLE estimate of heads-prob z given tosses x
def mle(x):
    return (x.sum()/(1+x*0).sum()).data.numpy()[0]

class MLE_Toss(nn.Module):
    def __init__(self, z=0.5):
        '''
        Compute ELBO using Monte Carlo with 'ns' samples
        '''
        super(MLE_Toss, self).__init__()

        self.z  =  Variable(t.FloatTensor([z]), requires_grad = True)


    def log_likelihood(self, x):
        '''
        mean (over tosses x_i ) of log( p(x_i | z) )
        '''
        p = t.sigmoid(self.z).repeat(len(x),1)
        return (x * t.log(p) + (1 - x) * t.log(1 - p)).mean()

    def parameters(self):
        '''
        Normally don't need this but this is a non-traditional network
        so we override parameters() so it returns the two params to optimize
        '''
        return [self.z]

    def forward(self, x):
        return -self.log_likelihood(x) # + self.kl()


nb = 100
bsiz = 100
epochs = 500
z = 0.3
x = gen_tosses(z, nb * bsiz)
mdl = MLE_Toss()

opt = optim.Adam(mdl.parameters(), lr = 0.01, betas = (0.8, 0.8),  eps=1e-7)
#opt = optim.SGD(mdl.parameters(), lr = 0.01)
#opt = optim.RMSprop(vt.parameters(), lr = 0.01)
# do one pass just to test

nll = mdl(x)
nll.backward()
opt.step()
[nll, mdl.z.grad, t.sigmoid(mdl.z)]

[p for p in mdl.parameters()]



for e in range(epochs):
    opt.zero_grad()
    nll = mdl(x)
    nll.backward()
    opt.step()
    if e % 10 == 0:
        print('nll= %8.3f z= %5.2f z.g= %5.2f  mle= %5.2f sigmoid_z= %5.2f' %
        (nll.data.numpy()[0],
        mdl.z.data.numpy()[0],
        mdl.z.grad.data.numpy()[0],
        mle(x),
        t.sigmoid(mdl.z).data.numpy()[0]
        ))

mdl.parameters()

