from network import *

# How much to move losing units in competitive learning
LOSER_ACTIVATION = .01

class CompetNetwork(Network):
    '''A competitive learning Network, with its own type of output Layer.'''

    def __init__(self, name, layers):
        '''Like Network, but sets supervised to False.'''
#        self.input_layer = layers[0]
        self.output_layers = layers[1:]
        Network.__init__(self, name, layers, supervised=False)
        self.competitive = True
        self.has_bias = False

    def step(self, pattern, train=True, show_act=True, lr=None):
        '''Run the network on one pattern, return the error and the winner (as a set).'''
        error = 0.0
        winners = set([])
        self.layers[0].clamp(pattern)
        for l in self.output_layers:
            # Tuple: (distance to winner, winner)
            err_win = l.update()
            error += err_win[0]
            winners.add(err_win[1])
            if train:
                l.learn()
        return error, winners

class CompetLayer(Layer):
    '''Output Layer for a CompetNetwork.'''

    def __init__(self, name, dimensions=(10, 10), lr=.1):
        Layer.__init__(self, name, size=0, weight_range=1.0, neg_weights=False, has_bias=False,
                       labels=[], lr=lr, linear=False, dimensions=dimensions)
        
    def loser_activation(self, winner, u):
        '''Activation for a loser; overridden in KohLayer.'''
        return LOSER_ACTIVATION

    def unit_distance(self, index):
        '''Distance between the unit with index and a clamped input pattern.'''
        return distance2(self.weights[index], self.input_layer.activations)

    def update(self):
        '''Update the network by finding the winner, giving it activation 1.0, others LOSER_ACTIVATION.'''
        winner, neg_error = argmax2(range(self.size), lambda x: -self.unit_distance(x))
        for u in range(self.size):
            if u == winner:
                self.activations[u] = 1.0
            else:
                self.activations[u] = self.loser_activation(winner, u)
        return math.sqrt(-neg_error), winner

    def learn(self):
        '''Update weights into all units, scaled by proportion of distance from input pattern.'''
        for j in range(self.size):
            act = self.activations[j]
            for i in range(self.input_layer.size):
                self.weights[j][i] += (self.input_layer.activations[i] - self.weights[j][i]) * act * self.lr

class KohLayer(CompetLayer):

    def __init__(self, name, dimensions=(10, 10), sigma=.8, decay=.9995, lr=.1):
        CompetLayer.__init__(self, name, dimensions=dimensions, lr=lr)
        self.sigma = sigma
        self.decay = decay

    def hidden_distance(self, u1, u2):
        '''Distance in the hidden layer between two units (only handles rectangles).'''
        return distance2(self.coords[u1], self.coords[u2]) / float(self.dimensions[0])

    def loser_activation(self, winner, u):
        '''The activation of losers is proportional to their distance from the winner.'''
        return self.neighborhood(winner, u)

    def neighborhood(self, u1, u2):
        '''Gaussian neighborhood function.'''
        return math.exp(-self.hidden_distance(u1, u2) / (2.0 * self.sigma * self.sigma))

    def learn(self):
        CompetLayer.learn(self)
        self.adjust_params()

    def adjust_params(self):
        self.sigma *= self.decay
        self.lr *= self.decay