from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.model_selection import KFold
import numpy as np
import scipy
from .utils import softmax, sigmoid, log_one_plus_exp

class RBM(BaseEstimator, TransformerMixin):
    """
    General Boltzmann Machine Model class,
    for vectorized computations.
    We assume: number of data points = each item is a column vector, and the data set is horizontally stacked
    We assume boolean values for all nodes! (not spin)
    """
    def __init__(self, n_visual, n_hidden):
        self.n_h = n_hidden
        self.n_v = n_visual
        self.W = None
        self.b_h = None
        self.b_v = None

    def initialize_weights(self, sigma=0.01):
        self.W = np.random.randn(self.n_v, self.n_h)*sigma
        #self.W = np.zeros((self.n_v, self.n_h))
        self.b_h = np.zeros(self.n_h).reshape(-1, 1)
        #self.b_h = np.random.randn((self.n_h)).reshape(-1, 1)
        self.b_v = np.zeros(self.n_v).reshape(-1, 1)

    def initialize_b_v(self, v):
        p = np.mean(v, axis=-1).reshape(-1, 1)
        p = np.minimum(p, 0.99)
        p = np.maximum(p, 0.01)
        self.b_v = np.log(p/(1-p))

    def energy(self, v, h):
        v = v.reshape(self.n_v, -1)
        h = h.reshape(self.n_h, -1)
        E = - np.einsum('...i,i...->...', np.dot(v.T, self.W), h).reshape(1, -1)
        E -= np.dot(self.b_h.T, h).reshape(1, -1) + np.dot(self.b_v.T, v).reshape(1, -1)
        return E

    def p_joint_vh(self, v, h):
        E = self.energy(v, h)
        return softmax(-E)

    def update_params(self, dw, dbh, dbv, alpha_w=0.001, alpha_b=0.01):
        self.W += alpha_w*dw
        self.b_h += alpha_b*dbh
        self.b_v += alpha_b*dbv

    def p_h_cond_v(self, v, with_b=True): # y_i = p(h_i = 1|v)
        v = v.reshape(self.n_v, -1)
        p = np.dot(self.W.T, v)
        if with_b: p = p + self.b_h # broadcast
        p = sigmoid(p)
        return p

    def p_v_cond_h(self, h, with_b=True): # y_i = p(v_i = 1|h)
        h = h.reshape(self.n_h, -1)
        p = np.dot(self.W, h)
        if with_b: p = p + self.b_v # broadcast
        p = sigmoid(p)
        return p

    def sample_binomial(self, p): # p = probability to be 1
        r = np.random.rand(*p.shape)
        return np.less(r, p).astype(float).reshape(*p.shape)

    def sample_p_x_cond_y(self, p_func, y):
        p = p_func(y)
        #print("probs = ",p)
        #print("p = ",p)
        return self.sample_binomial(p)

    def sample_p_h_cond_v(self, v):
        return self.sample_p_x_cond_y(self.p_h_cond_v, v)

    def sample_p_v_cond_h(self, h):
        return self.sample_p_x_cond_y(self.p_v_cond_h, h)

    def sample_vhv(self, v, k=1):
        #print("vbefore = ",v)
        vk = v.copy()
        for i in range(k):
            if i==0: # driven by data
                h = self.sample_p_h_cond_v(vk)
            else: #driven by reconstruction
                h = self.p_h_cond_v(vk)
#             print("hhh = ",h)
            #print("hhh = ", h.shape)
            #plot_n_images(h[:,:50], shape=(7,7))
            # always just probabilities
            vk = self.p_v_cond_h(h)
            #print("vvv = ", vk.shape)
            #plot_n_images(vk[:,:50], shape=(28,28))

            #print("vvv = ",vk)
        return vk


    def free_energy(self, v, mean=True):
        a = np.dot(self.W.T, v) + self.b_h # broadcast
        a = log_one_plus_exp(a)
        a = np.sum(a, axis=0).reshape(1, -1)
        a = - (np.dot(self.b_v.T, v) + a)
        if mean:
            return np.mean(a, axis=-1)
        else:
            return a


#     def log_p_v(self, v, x_min=0.01, x_max = 100):
#         a = np.dot(self.W.T, v) + self.b_h # broadcast
#         a = log_one_plus_exp(a)
#         a = np.sum(a, axis=0).reshape(1, -1)
#         return np.dot(self.b_v.T, v) + a




#     def log_p_v(self, v):
#         a = np.dot(self.W.T, v) + self.b_h # broadcast
#         print("dot",a)
#         a = np.cosh(a)
#         print("cosh",a)
#         a = np.log(a)
#         print("log",a)
#         a = np.sum(a, axis=0).reshape(1, -1)
#         print("sum",a)
#         return np.dot(self.b_v.T, v) + a

#     def log_p_v_CDk(self, v0, vk): # d log p_v / d theta
#         return self.log_p_v(v0) - self.log_p_v(vk)

    # http://deeplearning.net/tutorial/rbm.html#equation-free_energy_grad
    def derivatives_CDk(self, v0, vk):
        p_h_cond_v0 = self.p_h_cond_v(v0)
        p_h_cond_vk = self.p_h_cond_v(vk)
        dw = np.zeros((*self.W.shape, v0.shape[1]))
        for i in range(v0.shape[1]):
            dw[...,i] = np.outer(p_h_cond_v0[...,i], v0[...,i]).T - np.outer(p_h_cond_vk[...,i], vk[...,i]).T
            #print(p_h_cond_vk[...,i].shape, vk[...,i].T.shape)
            #print(dw[...,i])
        #print("P = ", p_h_cond_vk)
        #print("V = ", vk[...,i])
        dw = np.mean(dw, axis=-1)
        #print(dw)
        dbh = p_h_cond_v0 - p_h_cond_vk
        dbh = np.mean(dbh, axis=-1, keepdims=True)
        dbv = v0 - vk
        #print("v0 : ")
        #plot_n_images(v0[:,:50], shape=(28,28))
        #print("vk : ")
        #plot_n_images(vk[:,:50], shape=(28,28))
        #print("dbv (diff) : ")
        #plot_n_images(dbv[:,:50], shape=(28,28))

        dbv = np.mean(dbv, axis=-1, keepdims=True)
        #print("dbh = ",dbh.T)
        #print("dbv = ",dbv.T)
        return dw, dbh, dbv

    def step(self, v0, k=1, alpha_w=0.001, alpha_b=0.01):
        vk = self.sample_vhv(v0, k=k)
        dw, dbh, dbv = self.derivatives_CDk(v0, vk)
        rec_error = np.mean(np.abs(v0 - vk))
        #print("DW")
        #plot_bm_params(w=dw)
        self.update_params(dw, dbh, dbv, alpha_w=alpha_w, alpha_b=alpha_b)
        FE = self.free_energy(v0)
        return FE, rec_error

    def epoch(self, v0, k=1, alpha_w=0.001, alpha_b=0.01, batch_size=20):
        N = v0.shape[1]
        n_batch = (N-1) // batch_size + 1
        FE = np.zeros(n_batch)
        rec_error = np.zeros(n_batch)
        print("Epoch has {} batches".format(n_batch))
        if n_batch > 1:
            kf = KFold(n_splits=n_batch, shuffle=True)
            for i, (_, x_batch) in enumerate(kf.split(v0.T)):
                FE[i], rec_error[i] = self.step(v0[:,x_batch], k=k, alpha_w=alpha_w, alpha_b=alpha_b)
        else:
            FE[i], rec_error[i] = self.step(v0, k=k, alpha_w=alpha_w, alpha_b=alpha_b)
        return np.mean(FE), np.mean(rec_error)

    def reconstruct_hidden_units(self, with_b=True):
        h = np.eye(self.n_h) # activate one at a time
        return self.p_v_cond_h(h, with_b=with_b)
