"""
This is a scrappy first draft of a stimulus manager that will be built out
to incorporate arbitrary stimulus logic. For now you can subclass `Stim_Manager` and
redefine `next_stim`

TODO:
    Make this more general, for more than just sounds.
"""

import os
import pdb
from collections import deque
import numpy as np
import autopilot
from autopilot import prefs
if prefs.get('AGENT') and prefs.get('AGENT').upper() == 'PILOT':
    if 'AUDIO' in prefs.get('CONFIG') or prefs.get('AUDIOSERVER') is not None:
        from autopilot.stim.sound import sounds

def init_manager(stim):
    if 'manager' in stim.keys():
        if stim['manager'] in MANAGER_MAP.keys():
            manager = stim['manager']
            return MANAGER_MAP[manager](stim)
        else:
            Exception('Couldnt find stim manager of type: {}'.format(stim['type']))
    else:
        return Stim_Manager(stim)


class Stim_Manager(object):
    """
    Yield sounds according to some set of rules.

    Currently implemented:

    * correction trials - If a subject continually answers to one side incorrectly, keep
        the correct answer on the other side until they answer in that direction
    * bias correction - above some bias threshold, skew the correct answers to the less-responded side

    Attributes:
        stimuli (dict): Dictionary of instantiated stimuli like::

            {'L': [Tone1, Tone2, ...], 'R': [Tone3, Tone4, ...]}

        target ('L', 'R'): What is the correct port?
        distractor ('L', 'R'): What is the incorrect port?
        response ('L', 'R'): What was the last response?
        correct (0, 1): Was the last response correct?
        last_stim: What was the last stim? (one of `self.stimuli`)
        correction (bool): Are we doing correction trials?
        correction_trial (bool): Is this a correction trial?
        last_was_correction (bool): Was the last trial a correction trial?
        correction_pct (float): proportion of trials that are correction trials
        bias: False, or a bias correction mode.

    """

    # Metaclass for managing stimuli...
    def __init__(self, stim=None):
        """
        Args:
            stim (dict): Dictionary describing sound stimuli, in a format like::

                {
                'L': [{'type':'tone',...},{...}],
                'R': [{'type':'tone',...},{...}]
                }
        """
        self.stimuli = {}

        self.target = None  # What is the correct port?
        self.distractor = None  # What is the incorrect port
        self.response = None  # What was the last response?
        self.correct = 0  # Was the last response correct?
        self.last_stim = None  # What was the last stim?

        # Correction trials

        self.correction = False  # Are we doing correction trials
        self.correction_trial = False  # Is this a correction trial?
        self.last_was_correction = False  # Was the last trial a correction trial?
        self.correction_pct = 0.1  # proportion of trials that are correction trials

        # Bias correction
        self.bias = False  # or a bias correction mode

        # if we're being init'd as a superclass, no stim passed.
        if stim:

            # if we're doing sounds init them here
            # TODO: Make SoundManger subclass
            if 'sounds' in stim.keys():
                self.init_sounds(stim['sounds'])

    def do_correction(self, correction_pct = 0.5):
        """
        Called to set correction trials to True and correction percent.

        Args:
            correction_pct (float): Proportion of trials that should randomly be set to be correction trials.
        """
        self.correction = True
        self.correction_pct = correction_pct

    def do_bias(self, **kwargs):
        """
        Instantiate a :class:`.Bias_Correction` module

        Args:
            kwargs: parameters to initialize :class:`.Bias_Correction` with.
        """
        self.bias = Bias_Correction(**kwargs)

    def init_sounds(self, sound_dict):
        """
        Instantiate sound objects, using the 'type' value to choose an object from
        ``autopilot.get('sound')`` .

        Args:
            sound_dict (dict): a dictionary like::
                {
                'L': [{'type':'tone',...},{...}],
                'R': [{'type':'tone',...},{...}]
                }
        """
        # sounds should be
        # Iterate through sounds and load them to memory
        for k, v in sound_dict.items():
            # If multiple sounds on one side, v will be a list
            if isinstance(v, list):
                self.stimuli[k] = []
                for sound in v:
                    # We send the dict 'sound' to the function specified by 'type' and '
                    # ' as kwargs
                    self.stimuli[k].append(autopilot.get('sound', sound['type'])(**sound))
            # If not a list, a single sound
            else:
                self.stimuli[k] = [autopilot.get('sound', v['type'])(**v)]

    def set_triggers(self, trig_fn):
        """
        Give a callback function to all of our stimuli for when the stimulus ends.

        Note:
            Stimuli need a `set_trigger` method.

        Args:
            trig_fn (callable): A function to be given to stimuli via `set_trigger`
        """
        # set a callback function for when the stimulus ends
        for k, v in self.stimuli.items():
            for astim in v:
                astim.set_trigger(trig_fn)

    def make_punishment(self, type, duration):
        """
        Warning:
            Not Implemented

        Args:
            type:
            duration:
        """
        # types: timeout, noise
        # If we want a punishment sound...
        # if self.punish_sound:
        #     self.stimuli['punish'] = sounds.Noise(self.punish_dur)
        #     #change_to_green = lambda: self.hardware['LEDS']['C'].set_color([0, 255, 0])
        #     #self.stimuli['punish'].set_trigger(change_to_green)
        pass

    def play_punishment(self):
        """
        Warning:
            Not Implemented
        """
        pass

    def next_stim(self):
        """
        Compute and return the next stimulus

        If we are doing correction trials, compute that.

        Same thing with bias correction.

        Otherwise, randomly select a stimulus to present.

        Returns:
            ('L'/'R' Target, 'L'/'R' distractor, Stimulus to present)
        """
        # compute and return the next stim

        # first: if we're doing correction trials, compute that
        if self.correction:
            self.correction_trial = self.compute_correction()
            if self.correction_trial:
                return self.target, self.distractor, self.last_stim

        # otherwise we check for bias correction
        # it will return a threshold for random choice
        if self.bias:
            threshold = self.bias.next_bias()
        else:
            threshold = 0.5

        if np.random.rand()<threshold:
            self.target = 'L'
        else:
            self.target = 'R'

        if self.target == 'L':
            self.distractor = 'R'
        elif self.target == 'R':
            self.distractor = 'L'

        self.last_stim = np.random.choice(self.stimuli[self.target])

        return self.target, self.distractor, self.last_stim

    def compute_correction(self):
        """
        If `self.correction` is true, compute correction trial logic during
        `next_stim`.

        * If the last trial was a correction trial and the response to it wasn't correct, return True
        * If the last trial was a correction trial and the response was correct, return False
        * If the last trial as not a correction trial, but a randomly generated float is less than `correction_pct`, return True.

        Returns:
            bool: whether this trial should be a correction trial.

        """
        # if we are doing a correction trial this time return true

        # if this is the first trial, we can't do correction trials now can we
        if self.target is None:
            return False

        # if the last trial was a correction trial and we didn't get it correct,
        # then this is a correction trial too
        if (self.correction_trial or self.last_was_correction) and not self.correct:
            return True
        # if the last trial was a correction trial and we just corrected, no correction trial this time
        elif (self.correction_trial or self.last_was_correction) and self.correct:
            self.last_was_correction = False
            return False
        # if last trial was not a correction trial we spin  *to test* for one
        elif np.random.rand() < self.correction_pct:
            self.last_was_correction = True
            return False
        else:
            return False

    def update(self, response, correct):
        """
        At the end of a trial, update the status of our internal variables
        with the outcome of the trial.

        Args:
            response ('L', 'R'): How the subject responded
            correct (0, 1): Whether the response was correct.
        """
        self.response = response
        self.correct = correct
        if self.bias:
            self.bias.update(response, self.target)

    def end(self):
        """
        End all of our stim. Stim should have an `.end()` method of their own

        """

        for side, v in self.stimuli.items():
            for stim in v:
                try:
                    stim.end()
                except AttributeError:
                    print('stim does not have an end method! \n{}'.format(str(stim)))

class Proportional(Stim_Manager):
    """
    Present groups of stimuli with a particular frequency.

    Frequencies do not need to add up to 1, groups will be selected with the frequency
    (frequency)/(sum(frequencies)).

    Arguments:
        stim (dict): Dictionary with the structure:
            ::

                {'manager': 'proportional',
                 'type': 'sounds',
                 'groups': (
                     {'name':'group_name',
                      'frequency': 0.2,
                      'sounds':{
                          'L': [{Tone1_params}, {Tone2_params}...],
                          'R': [{Tone3_params}, {Tone4_params}...]
                      }
                    },
                    {'name':'second_group',
                      'frequency': 0.8,
                      'sounds':{
                          'L': [{Tone1_params}, {Tone2_params}...],
                          'R': [{Tone3_params}, {Tone4_params}...]
                      }
                    })
                }

    Attributes:
        stimuli (dict): A dictionary of stimuli organized into groups
        groups (dict): A dictionary mapping group names to frequencies

    """

    def __init__(self, stim):
        super(Proportional, self).__init__()

        self.stimuli = {}

        self.frequency_type = None

        if stim['type'] == 'sounds':
            if 'groups' in stim.keys():
                self.frequency_type = "within_groups"
                # top-level groups, choose group then choose side
                self.init_sounds_grouped(stim['groups'])
                self.store_groups(stim)
            else:
                self.frequency_type = "within_side"
                # second-level frequencies, side is chosen and then
                # probability from within a side
                self.init_sounds_individual(stim['sounds'])



    def init_sounds_grouped(self, sound_stim):
        """
        Instantiate sound objects similarly to :class:`.Stim_Manager`, just organizes them into groups.

        Args:
            sound_stim (tuple, list): an iterator like:
                ::

                    (
                     {'name':'group_name',
                      'frequency': 0.2,
                      'sounds': {
                          'L': [{Tone1_params}, {Tone2_params}...],
                          'R': [{Tone3_params}, {Tone4_params}...]
                        }
                    },
                    {'name':'second_group',
                      'frequency': 0.8,
                      'sounds':{
                          'L': [{Tone1_params}, {Tone2_params}...],
                          'R': [{Tone3_params}, {Tone4_params}...]
                      }
                    })
        """
        # Iterate through sounds and load them to memory
        self.stimuli = {}

        if isinstance(sound_stim, tuple) or isinstance(sound_stim, list):
            for group in sound_stim:
                group_name = group['name']

                # instantiate sounds
                self.stimuli[group_name] = {}
                for k, v in group['sounds'].items():
                    if isinstance(v, list):
                        self.stimuli[group_name][k] = []
                        for sound in v:
                            # We send the dict 'sound' to the function specified by 'type' and '
                            # ' as kwargs
                            self.stimuli[group_name][k].append(autopilot.get('sound', sound['type'])(**sound))
                    # If not a list, a single sound
                    else:
                        self.stimuli[group_name][k] = [autopilot.get('sound', v['type'])(**v)]


    def init_sounds_individual(self, sound_stim):
        """
        Initialize sounds with individually set presentation frequencies.

        .. todo::

            This method reflects the need for managers to have a unified schema,
            which will be built in a future release of Autopilot.

        Args:
            sound_stim (dict): Dictionary of {'side':[sound_params]} to generate sound stimuli

        Returns:

        """
        self.stim_freqs = {}
        for side, sound_params in sound_stim.items():
            self.stimuli[side] = []
            self.stim_freqs[side] = []
            if isinstance(sound_params, list):
                for sound in sound_params:
                    self.stimuli[side].append(autopilot.get('sound', sound['type'])(**sound))
                    self.stim_freqs[side].append(float(sound['management']['frequency']))
            else:
                self.stimuli[side].append(autopilot.get('sound', sound_params['type'])(**sound_params))
                self.stim_freqs[side].append(float(sound_params['management']['frequency']))

        # normalize frequencies within sides to sum to 1
        for side, freqs in self.stim_freqs.items():
            side_sum = np.sum(freqs)
            self.stim_freqs[side] = tuple([float(f)/side_sum for f in freqs])

    def store_groups(self, stim):
        """
        store groups and frequencies
        """
        self.groups = {}
        # also store (later) as tuples for choice b/c dicts are unordered
        group_names = []
        group_freqs = []
        for group in stim['groups']:
            group_name = group['name']
            frequency = float(group['frequency'])
            self.groups[group_name] = frequency
            group_names.append(group_name)
            group_freqs.append(frequency)

        self.group_names = tuple(group_names)
        group_freqs = np.array(group_freqs).astype(np.float)
        group_freqs = group_freqs/np.sum(group_freqs)

        self.group_freqs = tuple(group_freqs)



    def set_triggers(self, trig_fn):
        """
        Give a callback function to all of our stimuli for when the stimulus ends.

        Note:
            Stimuli need a `set_trigger` method.

        Args:
            trig_fn (callable): A function to be given to stimuli via `set_trigger`
        """
        # set a callback function for when the stimulus ends

        if self.frequency_type == "within_group":
            for _, group in self.stimuli.items():
                    for _, v in group.items():
                        for astim in v:
                            astim.set_trigger(trig_fn)
        elif self.frequency_type == "within_side":
            for side, sound_group in self.stimuli.items():
                for astim in sound_group:
                    astim.set_trigger(trig_fn)

        else:
            ValueError('Dont know how to set triggers')


    def next_stim(self):
        """
        Compute and return the next stimulus

        If we are doing correction trials, compute that.

        Same thing with bias correction.

        Otherwise, randomly select a stimulus to present, weighted by its group frequency.

        Returns:
            ('L'/'R' Target, 'L'/'R' distractor, Stimulus to present)
        """
        # compute and return the next stim

        # first: if we're doing correction trials, compute that
        if self.correction:
            self.correction_trial = self.compute_correction()
            if self.correction_trial:
                return self.target, self.distractor, self.last_stim

        # otherwise we check for bias correction
        # it will return a threshold for random choice
        if self.bias:
            threshold = self.bias.next_bias()
        else:
            threshold = 0.5

        # choose side
        if np.random.rand()<threshold:
            self.target = 'L'
        else:
            self.target = 'R'

        if self.target == 'L':
            self.distractor = 'R'
        elif self.target == 'R':
            self.distractor = 'L'

        if self.frequency_type == "within_group":

            # pick a stimulus based on group frequency
            group = np.random.choice(self.group_names, p=self.group_freqs)
            # within that group pick a random stimulus
            self.last_stim = np.random.choice(self.stimuli[group][self.target])

        elif self.frequency_type == "within_side":
            self.last_stim = np.random.choice(self.stimuli[self.target],
                                              p=self.stim_freqs[self.target])
        else:
            ValueError('Dont know what freq type we are')

        return self.target, self.distractor, self.last_stim





class Bias_Correction(object):
    """
    Basic Bias correction module. Modifies the threshold of random stimulus
    choice based on history of biased responses.

    Attributes:
        responses (:class:`collections.deque`): History of prior responses
        targets (:class:`collections.deque`): History of prior targets.
    """

    def __init__(self, mode='thresholded_linear', thresh=.2, window=100):
        """
        Args:
            mode: One of the following:

                * `'thresholded linear'` : above some threshold, do linear bias correction
                    eg. if response rate 65% left, make correct be right 65% of the time

            thresh (float): threshold above chance, ie. 0.2 means has to be 70% biased in window
            window (int): number of trials to calculate bias over
        """
        self.mode = mode
        self.threshold = float(thresh)
        self.window = int(window)
        self.responses = deque(maxlen=self.window)
        self.targets = deque(maxlen=self.window)

    def next_bias(self):
        """
        Compute the next bias depending on `self.mode`

        Returns:
            float: Some threshold :class:`.Stim_Manager` uses to decide left vs right.

        """

        # compute the bias threshold
        if self.mode == 'thresholded_linear':
            return self.thresholded_linear()

    def thresholded_linear(self):
        """
        If we are above the threshold, linearly correct the rate of
        presentation to favor the rarely responded side.

        eg. if response rate 65% left, make correct be right 65% of the time

        Returns:
            float: 0.5-bias, where bias is the difference between the mean response and mean target.
        """
        # if we are above threshold, return the bias
        bias = np.mean(self.responses)-np.mean(self.targets)
        if np.abs(bias)>self.threshold:
            return 0.5-bias
        else:
            return 0.5


    def update(self, response, target):
        """
        Store some new response and target values

        Args:
            response ('R', 'L'): Which side the subject responded to
            target ('R', 'L'): The correct side.
        """
        if isinstance(response, str):
            if response == "R":
                response = 1.0
            elif response == "L":
                response = 0.0

        if isinstance(target, str):
            if target == "R":
                target = 1.0
            elif target == "L":
                target = 0.0

        self.responses.append(float(response))
        self.targets.append(float(target))

MANAGER_MAP = {
    'proportional': Proportional
}





# class Reward_Manager(object):
#     def __init__(self):
#         pass
