#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
$Id$

Classes and fuinctions to control network and classifier operation.
"""

# just run the network.

import os
import sys
import time
import csv
import cPickle
import logging
lg = logging.getLogger(os.path.basename(__file__))
lg.setLevel(logging.INFO) # set to level 5 (< logging.DEBUG) to obtain pre and post patterns.

import configobj
import numpy

#import NeuroTools.signals as nts

from network import AntennalLobe
from network import BeeBrain
#from utils import utility_funcs
#import init_sim

def usage():
    print "Usage: %s configfile sim_type"%os.path.basename(__file__)
    print "configfile: Path to network configuration file"
    print "sim_type: sim or hw"

from neuclar.network_utilities import *

class ALController(object):
    """
    Setup the antennal lobe network and control simulation, data presentation etc.
    """
    def __init__(self, pynn, config):
        lg.info('building network')
        self.pynn = pynn
        self.net = AntennalLobe(pynn, config)

    def run_network(self, duration):
        lg.info('starting simulation for %.1f ms'%duration)
        self.pynn.run(duration)

    def set_pattern(self, pattern):
        lg.info("setting pattern in the AL.")
        self.net.set_pattern(pattern)

    def set_pattern_batch(self, patterns, time_per_pattern):
        """
        set the driver spike rates such that all patterns are presented in
        sequence.
        """
        lg.info('setting pattern batch in the AL.')
        self.net.set_batch_pattern(patterns, time_per_pattern)

    def retrieve_spikes(self):
        """
        Retrieve the spikes produced in the network.
        """
        lg.info('retrieving spikes.')
        return self.net.retrieve_spikes()        

class BrainController(object):
    """
    set up the honeybee brain, control simulation, present stimuli, learn.
    """
    def __init__(self, pynn, config):
        lg.info('setting up bee brain.')
        self.pynn = pynn
        self.config = config
        self.stim = 0
        self.brain = BeeBrain(pynn, config)
        assert (config['simulation']['calib_AL'] == 'False'),\
               "This version does not support AL calibration."
        if self.config['network'].has_key('randomize_weights_pndriver'):
            # driver -> PNs
            std = config['network'].as_float('randomize_weights_pndriver')
            if std > 0.:
                orig_weight = self.config['network'].as_float('w_driver_PN')
                self.brain.AL.randomize_pndriver_weights(orig_weight, std)
        if self.config['network'].has_key('randomize_weights_lndriver'):
            # PNs -> LNs
            std = config['network'].as_float('randomize_weights_lndriver')
            if std > 0.:
                orig_weight = config['network'].as_float('w_PN_LN')
                self.brain.AL.randomize_lndriver_weights(orig_weight, std)
        self.brain.AL.setup_lateral_inhibition_from_config()

    def set_pattern(self, pattern):
        """
        set the given input pattern as stimulus in the AL.
        """
        lg.info('setting stimulation pattern in the AL.')
        self.brain.AL.set_pattern(pattern)

    def set_pattern_batch(self, patterns, time_per_pattern):
        """
        set the given input pattern as stimulus in the AL.
        """
        lg.info('setting stimulation pattern in the AL.')
        self.brain.AL.set_pattern_batch(patterns, time_per_pattern)

    def run_network(self, duration):
        """
        run the network for duration ms.
        """
        lg.info('running the simulation for %.1f ms.'%duration)
        self.pynn.run(duration)
        lg.info('run completed.')
        
    def get_spikes(self, duration=None):
        """
        Retrieve spikes from all neurons after one stimulus.
        parameters: 
            duration - how far to look back for recorded spikes. If None, guess
                from config. 
        Returns dictionary with spike matrices.
        """
        if duration is None:
            duration = self.config['simulation'].as_float('duration')
        # get ORN spikes
        driverspikes = self.brain.AL.get_spikemat('drivers', 
                                                  not_older_than=duration)
        # get PN spikes
        pnspikes = self.brain.AL.get_spikemat('PNs', not_older_than=duration)
        # get LN spikes
        lnspikes = self.brain.AL.get_spikemat('LNs', not_older_than=duration)
        # get MB dec exc spikes 
        decexcspikes = self.brain.MBext.get_spikemat(pop='exc', 
                                                     not_older_than=duration)
        # get MB dec inh spikes 
        decinhspikes = self.brain.MBext.get_spikemat(pop='inh', 
                                                     not_older_than=duration)
        ret = {'drivers':driverspikes,
               'PNs': pnspikes,
               'LNs': lnspikes,
               'dec_exc': decexcspikes,
               'dec_inh': decinhspikes}
        return ret

    def test_pattern(self, pattern_tuple, class_ids='not used', timing_dict=None):
        """
        Present the pattern and determine the network's choice.

        Returns the number of spikes produced in each decision population.

        Parameters:
        pattern_tuple - pattern tuple as returned from PatternServer
                            (id, pattern, classlabel)
        class_ids - list of strings containing all possible class labels (not 
                        used but necessary in classifiers)
        timing_dict - dictionary in which times for 'run' and 'manage' will be
                    stored (for benchamrking).
        """
        start_time = time.time()
        lg.info('testing pattern.')
        self.stim +=1
        id = pattern_tuple[0]
        pattern = pattern_tuple[1]
        target = pattern_tuple[2]
        self.set_pattern(pattern)
        pat_creat_time = time.time()
        duration = self.config['simulation'].as_float('duration')
        self.run_network(duration)
        post_run_time = time.time()
        if self.pynn.__package__ == 'pyNN.hardware':
            t_back = None
        else:
            t_back = duration
        dn_spikecounts = self.brain.MBext.get_spikecountmat(
                                                        not_older_than=t_back)
        dec_pop_rates = numpy.mean(dn_spikecounts, axis=1)
        lg.info('pattern %s %s %s yielded response %s'%(id, str(pattern),