#!/usr/bin/env python
""" Reimplementation by Jens Kleesiek of 
A Sparse Generative Model of V1 Simple Cells with Intrinsic Plasticity
C. Weber and J. Triesch (2008) Neural Computation, 20, pp. 1261-84 
Usage:
First, download all files into one directory:
- helmholtz.py  (this file)
- retina.py     (for reading the image data patches)
- KTimage.py    (for exporting the weights and activations as pgm files for display)
- look.tcl      (to display those pgm files; look.tcl may be in any directory)
- data.tar.gz   (filtered images - preprocessing already done)
Second, tar xvzf data.tar.gz (creates a folder 'data' with the filtered images inside).
Third, create a directory /tmp/coco  for the weights and activations files.
Fourth, run: python helmholtz.py
After the weights and activation files are written, start:
./look.tcl a w 0 1  (left mouse click reloads the newest weights and activation files).
Fifth, wait ...
"""

import numpy as N
import retina as R
import KTimage as KT

print "Sparse Helmholtz machine for V1 responses\nusing Wake-Sleep Algorithm"

class helmholtz(object):
    def __init__(self, nvis, nhid, path):
        """ Data """
        self.__retinaSize = N.sqrt(nvis)
        self.__data = R.data(str(path), 481, 321, self.__retinaSize)

        """ Parameters """
        self.__nvis = nvis
        self.__nhid = nhid
        self.__alpha_bu = 0.1
        self.__alpha_td = 0.02

        """ Weights """ 
        self.__Wbu = N.random.uniform(-0.1,0.1,(self.__nhid,self.__nvis))
        self.__Wtd = N.random.uniform(-0.1,0.1,(self.__nvis,self.__nhid))

        """ Intrinsic Plasticity """
        self.__a = 2.5
        self.__b = -5.0
        self.__mu = 0.01
        self.__eta_a = 0.0001
        self.__eta_b = 0.0001

    def __fantasy_vector(self):
        """ generates binary vector via threshold """
        # generate uniform distributed numbers
        rnd = N.random.uniform(0,1,self.__nhid)
        fantasy = N.zeros(self.__nhid)
        val = 0.1
        fantasy[N.nonzero(rnd<self.__mu/val)[0]] = val
        return N.copy(fantasy)

    def __logistic_tf(self, y): 
        # activation hidden units (y)
        # eq 2.2
        tmp = self.__a * y + self.__b
        z = 1.0 / (1.0 + N.exp(-tmp))
        return N.copy(z)

    def __update_sparseness(self, y, z):
        # eq 2.4 and 2.5
        inv_mu = 1.0/self.__mu
        yz = y*z
        dA = 1.0/self.__a + y - 2*yz - inv_mu*yz + inv_mu*yz*z
        dB = 1.0 - 2*z - inv_mu*z + inv_mu*z*z
        self.__a += self.__eta_a * dA
        self.__b += self.__eta_b * dB

    def __normalize_rows(self, W):
        for i in range(N.shape(W)[0]):
            W[i] /= N.linalg.norm(W[i])

    def train(self, epoch):
        for itr in range(epoch):

            ## wake phase
            # get image patch
            input = 10.0 * (self.__data.getPatch()).reshape(self.__nvis,)
            # compute eq 2.1
            y = N.dot(self.__Wbu,input)
            # compute eq 2.2
            z = self.__logistic_tf(y) 
            # reconstruct
            x_rec = N.dot(self.__Wtd,z)
            if  normalised_mode:
                x_rec = x_rec / N.linalg.norm(x_rec) * N.linalg.norm(input)
            # compute eq 2.6, determine error x_hat
            x_hat = input - x_rec
            # compute eq 2.7, hebbian update
            self.__Wtd += self.__alpha_td * N.outer(x_hat,z)

            # update sparseness parameters eq. 2.4 & 2.5
            self.__update_sparseness(y,z)

            ## sleep phase
            # compute eq 2.8 & 2.9 using fantasy vectors
            # create sparse binary fantasy vector
            unicorn = self.__fantasy_vector()
            # project to input units
            x_tilde = N.dot(self.__Wtd,unicorn)
            # eq 2.1
            fi = self.__logistic_tf(N.dot(self.__Wbu, x_tilde))
            if  normalised_mode:
                fi = fi / N.linalg.norm(fi) * N.linalg.norm(unicorn)
            # reconstruction error
            z_hat = unicorn - fi
            # hebbian update
            self.__Wbu += self.__alpha_bu * N.outer(z_hat, x_tilde)

            # normalize weights
            if itr%10 == 0:
                self.__normalize_rows(self.__Wbu)
                if  normalised_mode:
                    self.__normalize_rows(self.__Wtd)

            if itr%5000 == 0:
                print 'Iteration %d'%itr, "m(unic)=", N.mean(unicorn), "m(fi)=%.3f"%N.mean(fi), "m(a(input))=%.3f"%N.mean(abs(input))
                KT.exporttiles (y, size_out_h, size_out_w, basedir+"obs_Y_1.pgm")            # y = Wbu * input     (wake)
                KT.exporttiles (z, size_out_h, size_out_w, basedir+"obs_Z_1.pgm")            # z = logistic(y)     (wake)
                KT.exporttiles (self.__a, size_out_h, size_out_w, basedir+"obs_A_1.pgm")     # a = IP param        (wake+sleep)
                KT.exporttiles (self.__b, size_out_h, size_out_w, basedir+"obs_B_1.pgm")     # b = IP param        (wake+sleep)
                KT.exporttiles (input, size_in_h, size_in_w, basedir+"obs_I_0.pgm")          # input               (wake)
                KT.exporttiles (x_rec, size_in_h, size_in_w, basedir+"obs_K_0.pgm")          # x_rec = Wtd * z     (wake)
                KT.exporttiles (self.__Wbu, size_in_h, size_in_w, basedir+"obs_W_1_0.pgm", size_out_h, size_out_w)                       # Wbu
                KT.exporttiles (N.transpose(N.copy(self.__Wtd)), size_in_h, size_in_w, basedir+"obs_W_0_1.pgm", size_out_h, size_out_w)  # Wtd
                KT.exporttiles (unicorn, size_out_h, size_out_w, basedir+"obs_J_1.pgm")      # unicorn             (sleep)
                KT.exporttiles (fi, size_out_h, size_out_w, basedir+"obs_K_1.pgm")           # fi = rec of unicorn (sleep)

basedir = "/tmp/coco/"
size_in_h, size_in_w = 8, 8
size_out_h, size_out_w = 10, 10
normalised_mode = 0
# normalised_mode = 0 (default):
#  only Wbu is normalised, Wtd adjusts by learning to reconstruct input by x_rec during wake
#  wake: input is given; z determined by IP; x_rec then fits to input
#  sleep: unicorn is given with IP mean (so resembles wake representation of input), fantasy-input should resemble input, fi should fit to unicorn (!?!)
#  note: size of Wtd will depend on amplitude of input! It may be convenient if Wbu and Wtd have comparable sizes so same learning rate can be used.
# normalised_mode = 1:
#  also Wtd is normalised; x_rec normalised to length of input; fi normalised to length of unicorn
#  hence, all weights are "fixed" and all reconstruction values fit in length to their targets - reduces scaling problems, easier, faster learning ...

if __name__ == "__main__":
    helm = helmholtz(size_in_h*size_in_w, size_out_h*size_out_w, 'data')
    helm.train(10000001)