# Princeton University licenses this file to You under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.  You may obtain a copy of the License at:
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.


# *************************************  AutoAssociativeLearningMechanism **********************************************

"""
.. _AutoAssociativeLearningMechanism_Overview:

Overview
--------

An AutoAssociativeLearningMechanism is a subclass `LearningMechanism <LearningMechanism>`, streamlined for use with a
`RecurrentTransferMechanism` to train its `recurrent_projection <RecurrentTransferMechanism.recurrent_projection>`.

.. _AutoAssociativeLearningMechanism_Creation:

Creating an AutoAssociativeLearningMechanism
--------------------------------------------

An AutoAssociativeLearningMechanism can be created directly by calling its constructor, but most commonly it is
created automatically when a RecurrentTransferMechanism is `configure for learning <Recurrent_Transfer_Learning>`,
(identified in its activity_source <AutoAssociativeLearningMechanism.activity_source>` attribute).

.. _AutoAssociativeLearningMechanism_Structure:

Structure
---------

An AutoAssociativeLearningMechanism is identical to a LearningMechanism in all respects except the following:

  * it has only a single *ACTIVATION_INPUT* `InputState`, that receives a `MappingProjection` from an `OutputState` of
    another `Mechanism` (identified by the `activity_source <AutoAssociativeLearningMechanism.activity_source>`,
    typically, the `primary OutputState <OutputState_Primary>` of a RecurrentTransferMechanism);

  * it has a single *LEARNING_SIGNAL* `OutputState` that sends a `LearningProjection` to the `matrix
    <AutoAssociativeProjection>` parameter of an 'AutoAssociativeProjection` (typically, the `recurrent_projection
    <RecurrentTransferMechanism.recurrent_projection>` of a RecurrentTransferMechanism), but not an
    *ERROR_SIGNAL* OutputState.

  * it has no `input_source <LearningMechanism.input_source>`, `output_source <LearningMechanism.output_source>`,
    or `error_source <LearningMechanism.error_source>` attributes;  instead, it has a single `activity_source`
    attribute that identifies the source of the activity vector used by the Mechanism's `function
    <AutoAssociativeLearningProjection.function>`.

  * it's `function <AutoAssociativeLearningMechanism.function>` takes as its `variable <Function_Base.variable>`
    a list or 1d np.array of numeric entries, corresponding in length to the AutoAssociativeLearningMechanism's
    *ACTIVATION_INPUT* InputState; and it returns a `learning_signal <LearningMechanism.learning_signal>`
    (a weight change matrix assigned to the Mechanism's *LEARNING_SIGNAL* OutputState), but not necessarily an
    `error_signal <LearningMechanism.error_signal>`.

  * its `learning_rate <AutoAssociativeLearningMechanism.learning_rate>` can be specified as a 1d or 2d array (or
    matrix) to scale the contribution made, respectively, by individual elements or connections among them,
    to the weight change matrix;  as with a standard `LearningMechanism`, a scalar can also be specified to scale
    the entire weight change matrix (see `learning_rate <AutoAssociativeLearningMechanism.learning_rate>` for
    additional details).

Execution
---------

An AutoAssociativeLearningMechanism executes in the same manner as standard `LearningMechanism`, however its execution
can be enabled or disabled by setting the the `learning_enabled <RecurrentTransferMechanism.learning_enabled>`
attribute of the `RecurrentTransferMechanism` with which it is associated (identified in its `activity_source
<AutoAssociativeLearningMechanism.attribute>`).

.. _AutoAssociativeLearningMechanism_Class_Reference:

Class Reference
---------------

"""

import numpy as np
import typecheck as tc

from psyneulink.components.component import parameter_keywords
from psyneulink.components.functions.function import Hebbian, ModulationParam, _is_modulation_param, is_function_type
from psyneulink.components.mechanisms.adaptive.learning.learningmechanism import ACTIVATION_INPUT, LearningMechanism
from psyneulink.components.mechanisms.processing.objectivemechanism import ObjectiveMechanism
from psyneulink.components.projections.pathway.mappingprojection import MappingProjection
from psyneulink.components.projections.projection import Projection_Base, _is_projection_spec, _validate_receiver, projection_keywords
from psyneulink.components.shellclasses import Projection
from psyneulink.globals.keywords import AUTOASSOCIATIVE_LEARNING_MECHANISM, CONTROL_PROJECTIONS, FUNCTION_PARAMS, INDEX, INITIALIZING, INPUT_STATES, LEARNING, LEARNING_PROJECTION, LEARNING_SIGNAL, LEARNING_SIGNALS, MAPPING_PROJECTION, MATRIX, NAME, OUTPUT_STATES, PROJECTION
from psyneulink.globals.preferences.componentpreferenceset import is_pref_set
from psyneulink.globals.preferences.preferenceset import PreferenceLevel
from psyneulink.globals.utilities import is_numeric, parameter_spec
from psyneulink.scheduling.timescale import CentralClock, TimeScale

__all__ = [
    'AutoAssociativeLearningMechanism', 'AutoAssociativeLearningMechanismError', 'DefaultTrainingMechanism',
    'input_state_names', 'output_state_names',
]

# Params:

parameter_keywords.update({LEARNING_PROJECTION, LEARNING})
projection_keywords.update({LEARNING_PROJECTION, LEARNING})

input_state_names =  [ACTIVATION_INPUT]
output_state_names = [LEARNING_SIGNAL]

DefaultTrainingMechanism = ObjectiveMechanism

class AutoAssociativeLearningMechanismError(Exception):
    def __init__(self, error_value):
        self.error_value = error_value

    def __str__(self):
        return repr(self.error_value)


class AutoAssociativeLearningMechanism(LearningMechanism):
    """
    AutoAssociativeLearningMechanism(              \
        variable,                                  \
        function=Hebbian,                          \
        learning_rate=None,                        \
        learning_signals=LEARNING_SIGNAL,          \
        modulation=ModulationParam.ADDITIVE,       \
        params=None,                               \
        name=None,                                 \
        prefs=None)

    Implements a `LearningMechanism` that modifies the `matrix <MappingProjection.matrix>` parameter of an
    `AutoAssociativeProjection` (typically the `recurrent_projection <RecurrentTransferMechanism.recurrent_projection>`
    of a `RecurrentTransferMechanism`).


    Arguments
    ---------

    variable : List or 2d np.array
        it must have a single item that corresponds to the value required by the AutoAssociativeLearningMechanism's
        `function <AutoAssociativeLearningMechanism..function>`;  it must each be compatible (in number and type)
        with the `value <InputState.value>` of the Mechanism's `InputState <LearningMechanism_InputStates>` (see
        `variable <AutoAssociativeLearningMechanism..variable>` for additional details).

    learning_signals : List[parameter of Projection, ParameterState, Projection, tuple[str, Projection] or dict]
        specifies the `matrix <AutoAssociativeProjection.matrix>` to be learned (see `learning_signals
        <LearningMechanism.learning_signals>` for details of specification).

    modulation : ModulationParam : ModulationParam.ADDITIVE
        specifies the default form of modulation used by the AutoAssociativeLearningMechanism's LearningSignals,
        unless they are `individually specified <LearningSignal_Specification>`.

    function : LearningFunction or function
        specifies the function used to calculate the AutoAssociativeLearningMechanism's `learning_signal
        <AutoAssociativeLearningMechanism.learning_signal>` attribute.  It must take as its **variable** argument a
        list or 1d array of numeric values (the "activity vector") and return a list, 2d np.array or np.matrix
        representing a square matrix with dimensions that equal the length of its variable (the "weight change
        matrix").

    learning_rate : float
        specifies the learning rate for the AutoAssociativeLearningMechanism. (see `learning_rate
        <AutoAssociativeLearningMechanism.learning_rate>` for details).

    params : Dict[param keyword, param value] : default None
        a `parameter dictionary <ParameterState_Specification>` that specifies the parameters for the
        Projection, its function, and/or a custom function and its parameters. By default, it contains an entry for
        the Projection's default `function <LearningProjection.function>` and parameter assignments.  Values specified
        for parameters in the dictionary override any assigned to those parameters in arguments of the constructor.

    name : str : default see `name <AutoAssociativeLearningMechanism.name>`
        specifies the name of the AutoAssociativeLearningMechanism.

    prefs : PreferenceSet or specification dict : default Mechanism.classPreferences
        specifies the `PreferenceSet` for the AutoAssociativeLearningMechanism; see `prefs <AutoAssociativeLearningMechanism.prefs>` for details.


    Attributes
    ----------

    COMMENT:
        componentType : LEARNING_MECHANISM
    COMMENT

    variable : 2d np.array
        has a single item, that serves as the template for the input required by the AutoAssociativeLearningMechanism's
        `function <AutoAssociativeLearningMechanism.function>`, corresponding to the `value
        <OutputState.value>` of the `activity_source <AutoAssociativeLearningMechanism.activity_source>`.

    activity_source : OutputState
        the `OutputState` that is the `sender <AutoAssociativeProjection.sender>` of the `AutoAssociativeProjection`
        that the Mechanism trains.

    input_states : ContentAddressableList[OutputState]
        has a single item, that contains the AutoAssociativeLearningMechanism's single *ACTIVATION_INPUT* `InputState`.

    primary_learned_projection : AutoAssociativeProjection
        the `Projection` with the `matrix <AutoAssociativeProjection.matrix>` parameter being trained by the
        AutoAssociativeLearningMechanism.  It is always the first Projection listed in the
        AutoAssociativeLearningMechanism's `learned_projections <AutoAssociativeLearningMechanism.learned_projections>`
        attribute.

    learned_projections : List[MappingProjection]
        all of the `AutoAssociativeProjections <AutoAssociativeProjection>` modified by the
        AutoAssociativeLearningMechanism;  the first item in the list is always the `primary_learned_projection
        <AutoAssociativeLearningMechanism.primary_learned_projection>`.

    function : LearningFunction or function : default Hebbian
        the function used to calculate the `learning_signal <AutoAssociativeLearningMechanism.learning_signal>`
        (assigned to the AutoAssociativeLearningMechanism's `LearningSignal(s) <LearningMechanism_LearningSignal>`).
        It's `variable <Function_Base.variable>` must be a list or 1d np.array of numeric entries, corresponding in
        length to the AutoAssociativeLearningMechanism's *ACTIVATION_INPUT* (`primary <InputState_Primary>`) InputState.

    learning_rate : float, 1d or 2d np.array, or np.matrix of numeric values : default None
        determines the learning rate used by the AutoAssociativeLearningMechanism's `function
        <AutoAssociativeLearningMechanism.function>` to scale the weight change matrix it returns. If it is a scalar, it is used to multiply the weight change matrix;  if it is a 2d array or matrix,
        it is used to Hadamard (elementwise) multiply the weight matrix (allowing the contribution of individual
        *connections* to be scaled);  if it is a 1d np.array, it is used to Hadamard (elementwise) multiply the input
        to the `function <AutoAssociativeLearningMechanism.function>` (i.e., the `value <InputState.value>` of the
        AutoAssociativeLearningMechanism's *ACTIVATION_INPUT* `InputState <AutoAssociativeLearningMechanism_Structure>`,
        allowing the contribution of individual *units* to be scaled). If specified, the value supersedes the
        learning_rate assigned to any `Process` or `System` to which the AutoAssociativeLearningMechanism belongs.
        If it is `None`, then the `learning_rate <Process.learning_rate>` specified for the Process to which the
        AutoAssociativeLearningMechanism belongs belongs is used;  and, if that is `None`, then the `learning_rate
        <System.learning_rate>` for the System to which it belongs is used. If all are `None`, then the
        `default_learning_rate <LearningFunction.default_learning_rate>` for the `function
        <AutoAssociativeLearningMechanism.function>` is used (see `learning_rate <LearningMechanism_Learning_Rate>`
        for additional details).

    learning_signal : 2d ndarray or matrix of numeric values
        the value returned by `function <AutoAssociativeLearningMechanism.function>`, that specifies
        the changes to the weights of the `matrix <AutoAssociativeProjection.matrix>` parameter for the
        AutoAssociativeLearningMechanism's`learned_projections <AutoAssociativeLearningMechanism.learned_projections>`;
        It is assigned as the value of the AutoAssociativeLearningMechanism's `LearningSignal(s)
        <LearningMechanism_LearningSignal>` and, in turn, its `LearningProjection(s) <LearningProjection>`.

    learning_signals : List[LearningSignal]
        list of all of the `LearningSignals <LearningSignal>` for the AutoAssociativeLearningMechanism, each of which
        sends one or more `LearningProjections <LearningProjection>` to the `ParameterState(s) <ParameterState>` for
        the `matrix <AutoAssociativeProjection.matrix>` parameter of the `AutoAssociativeProjection(s)
        <AutoAssociativeProjection>` trained by the AutoAssociativeLearningMechanism.  Although in most instances an
        AutoAssociativeLearningMechanism is used to train a single AutoAssociativeProjection, like a standard
        `LearningMechanism`, it can be assigned additional LearningSignals and/or LearningProjections to train
        additional ones;  in such cases, the `value <LearningSignal>` for all of the LearningSignals is the
        the same:  the AutoAssociativeLearningMechanism's `learning_signal
        <AutoAssociativeLearningMechanism.learning_signal>` attribute, based on its `activity_source
        <AutoAssociativeLearningMechanism.activity_source>`.  Since LearningSignals are `OutputStates
        <OutputState>`, they are also listed in the AutoAssociativeLearningMechanism's `output_states
        <AutoAssociativeLearningMechanism.output_states>` attribute.

    learning_projections : List[LearningProjection]
        list of all of the LearningProjections <LearningProjection>` from the AutoAssociativeLearningMechanism, listed
        in the order of the `LearningSignals <LearningSignal>` to which they belong (that is, in the order they are
        listed in the `learning_signals <AutoAssociativeLearningMechanism.learning_signals>` attribute).

    output_states : ContentAddressableList[OutputState]
        list of the AutoAssociativeLearningMechanism's `OutputStates <OutputState>`, beginning with its
        `learning_signals <AutoAssociativeLearningMechanism.learning_signals>`, and followed by any additional
        (user-specified) `OutputStates <OutputState>`.

    output_values : 2d np.array
        the first item is the `value <OutputState.value>` of the LearningMechanism's `learning_signal
        <AutoAssociativeLearningMechanism.learning_signal>`, followed by the `value <OutputState.value>`\\s
        of any additional (user-specified) OutputStates.

    modulation : ModulationParam
        the default form of modulation used by the AutoAssociativeLearningMechanism's `LearningSignal(s)
        <LearningMechanism_LearningSignal>`, unless they are `individually specified <LearningSignal_Specification>`.

    name : str
        the name of the AutoAssociativeLearningMechanism; if it is not specified in the **name** argument of the
        constructor, a default is assigned by MechanismRegistry (see `Naming` for conventions used for default and
        duplicate names).

    prefs : PreferenceSet or specification dict
        the `PreferenceSet` for the AutoAssociativeLearningMechanism; if it is not specified in the **prefs** argument
        of the constructor, a default is assigned using `classPreferences` defined in __init__.py (see
        `PreferenceSet <LINK>` for details).
    """

    componentType = AUTOASSOCIATIVE_LEARNING_MECHANISM
    className = componentType
    suffix = " " + className

    classPreferenceLevel = PreferenceLevel.TYPE

    # ClassDefaults.variable = None

    paramClassDefaults = Projection_Base.paramClassDefaults.copy()
    paramClassDefaults.update({
        CONTROL_PROJECTIONS: None,
        INPUT_STATES:input_state_names,
        OUTPUT_STATES:[{NAME:LEARNING_SIGNAL,  # NOTE: This is the default, but is overridden by any LearningSignal arg
                        INDEX:0}
                       ]})

    @tc.typecheck
    def __init__(self,
                 variable:tc.any(list, np.ndarray),
                 size=None,
                 function:is_function_type=Hebbian,
                 learning_signals:tc.optional(list) = None,
                 modulation:tc.optional(_is_modulation_param)=ModulationParam.ADDITIVE,
                 learning_rate:tc.optional(parameter_spec)=None,
                 params=None,
                 name=None,
                 prefs:is_pref_set=None,
                 context=None):

        # Assign args to params and functionParams dicts (kwConstants must == arg names)
        params = self._assign_args_to_param_dicts(function=function,
                                                  learning_signals=learning_signals,
                                                  params=params)

        # # USE FOR IMPLEMENTATION OF deferred_init()
        # # Store args for deferred initialization
        # self.init_args = locals().copy()
        # self.init_args['context'] = self
        # self.init_args['name'] = name

        # # Flag for deferred initialization
        # self.init_status = InitStatus.DEFERRED_INITIALIZATION

        self._learning_rate = learning_rate

        super().__init__(variable=variable,
                         size=size,
                         function=function,
                         modulation=modulation,
                         params=params,
                         name=name,
                         prefs=prefs,
                         context=self)

    def _validate_variable(self, variable, context=None):
        """Validate that variable has only one item: activation_input.
        """

        # Skip LearningMechanism._validate_variable in call to super(), as it requires variable to have 3 items
        variable = self._update_variable(super(LearningMechanism, self)._validate_variable(variable, context))

        # MODIFIED 9/22/17 NEW: [HACK]
        if np.array(np.squeeze(variable)).ndim != 1 or not is_numeric(variable):
        # MODIFIED 9/22/17 END
            raise AutoAssociativeLearningMechanismError("Variable for {} ({}) must be "
                                                        "a list or 1d np.array containing only numbers".
                                                        format(self.name, variable))
        return variable

    def _execute(self,
                variable=None,
                runtime_params=None,
                clock=CentralClock,
                time_scale = TimeScale.TRIAL,
                context=None):
        """Execute AutoAssociativeLearningMechanism. function and return learning_signal

        :return: (2D np.array) self.learning_signal
        """

        # COMPUTE LEARNING SIGNAL (note that function is assumed to return only one value)
        self.learning_signal = self.function(variable=variable,
                                             params=runtime_params,
                                             context=context)

        if not INITIALIZING in context and self.reportOutputPref:
            print("\n{} weight change matrix: \n{}\n".format(self.name, self.learning_signal))

        self.value = [self.learning_signal]
        return self.value

    @property
    def activity_source(self):
        return self.input_state.path_afferents[0].sender.owner
