"""Represent interactions between context and item layers."""

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from cymr import operations


def expand_param(param, size):
    """
    Expand a scalar parameter to array format.

    Parameters
    ----------
    param : float or numpy.ndarray
        Parameter to expand.

    size : iterable or numpy.ndarray
        Size of the expanded parameter.

    Returns
    -------
    param : numpy.ndarray
        Expanded parameter.
    """
    size = np.asarray(size)
    if not isinstance(param, np.ndarray):
        # expand scalar to full array
        param = np.tile(param, size).astype(float)
    elif size.shape and param.ndim < len(size):
        # expand array to have to correct number of dimensions
        param = np.expand_dims(param, axis=1)
        if param.shape != tuple(size):
            # expand singleton dimensions as needed
            rep = np.ones(size.shape, dtype=int)
            for i, n in enumerate(param.shape):
                if n == 1:
                    rep[i] = size[i]
                elif n != size[i] and n < size[i]:
                    raise ValueError('Cannot expand parameter.')
            param = np.tile(param, rep).astype(float)
    return param


def prepare_study_param(n_item, n_sub, B, Lfc, Lcf, distract_B=None, retention_B=None):
    """Prepare parameters for simulating a study phase."""
    B = expand_param(B, (n_item, n_sub))
    Lfc = expand_param(Lfc, (n_item, n_sub))
    Lcf = expand_param(Lcf, (n_item, n_sub))
    param = {'B': B, 'Lfc': Lfc, 'Lcf': Lcf, 'distract_B': None}
    if distract_B is not None:
        distract_B = expand_param(distract_B, (n_item, n_sub))
        if isinstance(retention_B, np.ndarray) and retention_B.shape[0] > 1:
            # if there is a value for each item, use the last one
            retention_B = retention_B[[-1]]
        retention_B = expand_param(retention_B, (1, n_sub))
        param['distract_B'] = np.concat([distract_B, retention_B], 0)
    return param


def prepare_recall_param(n_item, n_sub, B, T, amin):
    """Prepare parameters for simulating a recall phase."""
    B = expand_param(B, (n_item, n_sub))
    if T < amin:
        T = amin
    param = {'B': B, 'T': T}
    return param


def sample_response_lba(A, b, v, s, tau):
    """Sample a response and response time."""
    while True:
        k = stats.uniform.rvs(0, scale=A)
        d = stats.norm.rvs(loc=v, scale=s)
        t = tau + (b - k) / d
        if np.any(d > 0):
            break
    t[d <= 0] = np.nan
    response = np.nanargmin(t)
    rt = np.nanmin(t)
    return int(response), rt


def init_plot(**kwargs):
    fig = plt.figure(constrained_layout=True, **kwargs)
    gs = GridSpec(10, 5, figure=fig)
    ax = {
        'c': fig.add_subplot(gs[0, 1:4]),
        'c_in': fig.add_subplot(gs[1, 1:4]),
        'f_in': fig.add_subplot(gs[8, 1:4]),
        'f': fig.add_subplot(gs[9, 1:4]),
        'w_fc_pre': fig.add_subplot(gs[2:8, 0]),
        'w_fc_exp': fig.add_subplot(gs[2:8, 1]),
        'w_ff_pre': fig.add_subplot(gs[2:8, 2]),
        'w_cf_exp': fig.add_subplot(gs[2:8, 3]),
        'w_cf_pre': fig.add_subplot(gs[2:8, 4]),
    }
    return fig, ax


class LayerIndex(object):
    """
    Representation of one layer of a network.

    Parameters
    ----------
    layer_segments : dict of (str: dict of (str: int))
        Size of each named segment for each sublayer.

    Attributes
    ----------
    size : int
        Number of units in the layer.

    size_sublayer : dict of (str: int)
        Number of units in each sublayer.

    size_segment : dict of (str: dict of (str: int))
        Size of each segment in each sublayer.

    sublayer : dict of (str: numpy.ndarray)
        Start and stop units of each sublayer.

    segment : dict of (str: dict of (str: numpy.ndarray)
        Start and stop units of each segment.
    """

    def __init__(self, layer_segments):
        self.size = 0
        self.size_sublayer = {}
        self.size_segment = layer_segments.copy()
        self.sublayer = {}
        self.segment = {}
        for sub, segs in layer_segments.items():
            self.segment[sub] = {}
            self.size_sublayer[sub] = 0
            start = self.size
            for seg, s in segs.items():
                self.segment[sub][seg] = np.array(
                    [self.size, self.size + s], dtype=np.dtype('i')
                )
                self.size += s
                self.size_sublayer[sub] += s
            self.sublayer[sub] = np.array(
                [start, start + self.size_sublayer[sub]], dtype=np.dtype('i')
            )

    def __repr__(self):
        s = ''
        for sublayer, segments in self.size_segment.items():
            size_sublayer = self.size_sublayer[sublayer]
            s += f'{sublayer}: {size_sublayer} units\n'
            for segment, size_segment in segments.items():
                s += f'    {segment}: {size_segment} units\n'
        return s

    def copy(self):
        return self.__init__(self.size_segment)

    def get_sublayer(self, sublayer):
        return self.sublayer[sublayer]

    def get_segment(self, sublayer, segment):
        return self.segment[sublayer][segment]

    def get_slice(self, sublayer, segment):
        ind = self.get_segment(sublayer, segment)
        return slice(ind[0], ind[1])

    def get_unit(self, sublayer, segment, index):
        return self.segment[sublayer][segment][0] + index


class Network(object):
    """
    Representation of interacting item and context layers.

    Parameters
    ----------
    f_segment : dict of str: (dict of str: int)
        For each item sublayer, the number of units in each segment.

    c_segment : dict of str: (dict of str: int)
        For each context sublayer, the number of units in each segment.

    Attributes
    ----------
    f_segment : dict of str: (dict of str: int)
        Number of item units for each segment.

    c_segment : dict of str: (dict of str: int)
        Number of context units for each segment.

    f_ind : cymr.network.LayerIndex
        Index of units in the item layer.

    c_ind : cymr.network.LayerIndex
        Index of units in the context layer.

    n_f : int
        Total number of item units.

    n_c : int
        Total number of context units.

    f : numpy.array
        Item layer vector.

    c : numpy.array
        Context layer vector.

    c_in : numpy.array
        Current input to context.

    w_fc_pre : numpy.array
        Pre-experimental weights connecting f to c.

    w_fc_exp : numpy.array
        Weights learned during the experiment connecting f to c.

    w_cf_pre : numpy.array
        Pre-experimental weights connecting c to f.

    w_cf_exp : numpy.array
        Weights learned during the experiment connecting c to f.

    w_ff_pre : numpy.array
        Pre-experimental weights connecting f to f.

    w_ff_exp : numpy.array
        Weights learned during the experiment connecting f to f.
    """

    def __init__(self, f_segment, c_segment):
        self.f_segment = f_segment.copy()
        self.c_segment = c_segment.copy()
        self.f_sublayers = list(self.f_segment.keys())
        self.c_sublayers = list(self.c_segment.keys())
        self.f_ind = LayerIndex(self.f_segment)
        self.c_ind = LayerIndex(self.c_segment)
        n_f = self.f_ind.size
        n_c = self.c_ind.size
        self.n_f = n_f
        self.n_c = n_c
        self.f = np.zeros(n_f)
        self.f_in = np.zeros(n_f)
        self.match = np.zeros(n_f)
        self.c = np.zeros(n_c)
        self.c_in = np.zeros(n_c)
        self.w_fc_pre = np.zeros((n_f, n_c))
        self.w_fc_exp = np.zeros((n_f, n_c))
        self.w_cf_pre = np.zeros((n_f, n_c))
        self.w_cf_exp = np.zeros((n_f, n_c))
        self.w_ff_pre = np.zeros((n_f, n_f))
        self.w_ff_exp = np.zeros((n_f, n_f))

    def __repr__(self):
        s = f'f:\n{self.f_ind}\nc:\n{self.c_ind}'
        return s

    def reset(self):
        """Reset network weights and activations to zero."""
        self.f[:] = 0
        self.f_in[:] = 0
        self.match[:] = 0
        self.c[:] = 0
        self.c_in[:] = 0
        self.w_fc_exp[:] = 0
        self.w_fc_pre[:] = 0
        self.w_cf_exp[:] = 0
        self.w_cf_pre[:] = 0
        self.w_ff_exp[:] = 0
        self.w_ff_pre[:] = 0

    def copy(self, include=None, exclude=None):
        """
        Copy the network to a new network object.

        Parameters
        ----------
        include : list of str, optional
            List of fields to include. Default is a standard list of
            fields.

        exclude : list of str, optional
            List of fields to exclude. Overrides values for include.

        Returns
        -------
        net : cymr.Network
            Network with the same segments, weights, and activations.
        """
        # all attributes that contain an array
        fields = [
            'f',
            'f_in',
            'match',
            'c',
            'c_in',
            'w_fc_exp',
            'w_fc_pre',
            'w_cf_exp',
            'w_cf_pre',
            'w_ff_exp',
            'w_ff_pre',
        ]

        # set the fields to include in the copy (default is all)
        if include is None:
            include = fields
        if exclude is not None:
            include = [i for i in include if i not in exclude]

        # copy included attributes
        net = Network(self.f_segment, self.c_segment)
        for field in fields:
            if field in include:
                setattr(net, field, getattr(self, field).copy())
            else:
                delattr(net, field)
        return net

    def get_sublayer(self, layer, sublayer):
        """
        Get indices for a sublayer.

        Parameters
        ----------
        layer : str
            Layer to access.

        sublayer : str
            Sublayer to access.

        Returns
        -------
        numpy.ndarray
            Start and stop indices for the sublayer.
        """
        if layer == 'f':
            ind = self.f_ind.get_sublayer(sublayer)
        elif layer == 'c':
            ind = self.c_ind.get_sublayer(sublayer)
        else:
            raise ValueError(f'Invalid layer: {layer}')
        return ind

    def get_sublayers(self, layer, sublayers):
        """
        Get an array of indices for multiple sublayers.

        Parameters
        ----------
        layer : str
            Layer to access.

        sublayers : list of str
            Sublayers to access.

        Returns
        -------
        numpy.array
            [sublayers x 2] array of start and stop indices for each
            sublayer.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        ind_list = []
        for sublayer in sublayers:
            ind_list.append(self.get_sublayer(layer, sublayer))
        return np.array(ind_list, dtype=np.dtype('i'))

    def get_segment(self, layer, sublayer, segment):
        """
        Get indices for a segment.

        Parameters
        ----------
        layer : {'f', 'c'}
            Layer to access.

        sublayer : str
            Sublayer to access.

        segment : str
            Segment to access.

        Returns
        -------
        ind : numpy.array
            Segment indices.
        """
        if layer == 'f':
            ind = self.f_ind.get_segment(sublayer, segment)
        elif layer == 'c':
            ind = self.c_ind.get_segment(sublayer, segment)
        else:
            raise ValueError(f'Invalid layer: {layer}')
        return ind

    def get_unit(self, layer, sublayer, segment, unit):
        """
        Get indices for a unit.

        Parameters
        ----------
        layer : {'f', 'c'}
            Layer to access.

        sublayer : str
            Sublayer to access.

        segment : str
            Segment to access.

        unit : int
            Index of the unit within the segment.

        Returns
        -------
        ind : int
            Absolute index of the unit within the layer.
        """
        if layer == 'f':
            ind = self.f_ind.get_unit(sublayer, segment, unit)
        elif layer == 'c':
            ind = self.c_ind.get_unit(sublayer, segment, unit)
        else:
            raise ValueError(f'Invalid layer: {layer}')
        return ind

    def get_slice(self, layer, sublayer, segment):
        """
        Get a slice for a segment.

        Parameters
        ----------
        layer : {'f', 'c'}
            Layer to access.

        sublayer : str
            Sublayer to access.

        segment : str
            Segment to access.

        Returns
        -------
        slice
            Slice for indexing the segment.
        """
        if layer == 'f':
            index = self.f_ind
        elif layer == 'c':
            index = self.c_ind
        else:
            raise ValueError(f'Invalid layer: {layer}')
        return index.get_slice(sublayer, segment)

    def get_region(self, f_segment, c_segment):
        """
        Return slices for a region.

        Parameters
        ----------
        f_segment : tuple of (str, str)
            Sublayer and segment of the item layer.

        c_segment : tuple of (str, str)
            Sublayer and segment of the context layer.

        Returns
        -------
        f_ind : slice
            Span of the region in the item dimension.

        c_ind : slice
            Span of the region in the context dimension.
        """
        f_ind = self.f_ind.get_slice(*f_segment)
        c_ind = self.c_ind.get_slice(*c_segment)
        return f_ind, c_ind

    def add_pre_weights(
        self, connect, f_segment, c_segment, weights, slope=1, intercept=0
    ):
        """
        Add pre-experimental weights to a network.

        Parameters
        ----------
        connect : {'fc', 'cf'}
            Connections to add weights to.

        f_segment : tuple of str, str
            Sublayer and segment to use for the item layer.

        c_segment : tuple of str, str
            Sublayer and segment to use for the context layer.

        weights : numpy.array
            Items x context array of weights.

        slope : double, optional
            Slope to multiply weights by before adding.

        intercept : double, optional
            Intercept to add to weights.
        """
        scaled = intercept + slope * weights
        if connect == 'ff':
            f_slice = slice(*tuple(self.f_ind.get_segment(*f_segment)))
            c_slice = None
        else:
            f_slice, c_slice = self.get_region(f_segment, c_segment)
        if connect == 'fc':
            self.w_fc_pre[f_slice, c_slice] = scaled
        elif connect == 'cf':
            self.w_cf_pre[f_slice, c_slice] = scaled
        elif connect == 'ff':
            self.w_ff_pre[f_slice, f_slice] = scaled
        else:
            raise ValueError(f'Invalid connection type: {connect}')

    def update(self, item, sublayers):
        """
        Update context completely with input from the item layer.

        Rather than integrating input into context, this replaces the
        current state of context with the input.

        Parameters
        ----------
        item : tuple of str, str, int
            Sublayer, segment, and unit of the item to present.

        sublayers : str of list of str
            Sublayer(s) of context to update.
        """
        f_ind = self.get_unit('f', *item)
        c_ind = self.get_sublayers('c', sublayers)
        B = np.ones(len(sublayers))
        operations.integrate(
            self.w_fc_exp, self.w_fc_pre, self.c, self.c_in, self.f, f_ind, c_ind, B
        )

    def integrate(self, item, sublayers, B):
        """
        Integrate input from the item layer into context.

        Parameters
        ----------
        item : tuple of str, str, int
            Sublayer, segment, and unit of the item to present.

        sublayers : str or list of str
            Sublayer(s) of context to update.

        B : float or numpy.array
            Integration scaling factor for each sublayer; higher values
            update context to more strongly reflect the input.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        B = expand_param(B, np.array(len(sublayers)))
        f_ind = self.get_unit('f', *item)
        c_ind = self.get_sublayers('c', sublayers)
        operations.integrate(
            self.w_fc_exp, self.w_fc_pre, self.c, self.c_in, self.f, f_ind, c_ind, B
        )

    def present(self, item, sublayers, B, Lfc=0, Lcf=0):
        """
        Present an item and learn context-item associations.

        Parameters
        ----------
        item : tuple of str, str, int
            Sublayer, segment, and unit of the item to present.

        sublayers : str or list of str
            Sublayer(s) of context to update.

        B : float or numpy.ndarray
            Integration scaling factors; higher values update context
            to more strongly reflect the input.

        Lfc : float or numpy.ndarray, optional
            Learning rates for item to context associations.

        Lcf : float or numpy.ndarray, optional
            Learning rates for context to item associations.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        n_sub = np.array(len(sublayers))
        B = expand_param(B, n_sub)
        Lfc = expand_param(Lfc, n_sub)
        Lcf = expand_param(Lcf, n_sub)
        f_ind = self.get_unit('f', *item)
        c_ind = self.get_sublayers('c', sublayers)
        operations.present(
            self.w_fc_exp,
            self.w_fc_pre,
            self.w_cf_exp,
            self.c,
            self.c_in,
            self.f,
            f_ind,
            c_ind,
            B,
            Lfc,
            Lcf,
        )

    def learn(self, connect, item, sublayers, L):
        """
        Learn an association between the item and context layers.

        Parameters
        ----------
        connect : {'fc', 'cf'}
            Connection matrix to update.

        item : tuple of str, str, int
            Sublayer, segment, and unit of the item to present.

        sublayers : str
            Sublayers of context to update.

        L : double
            Learning rate.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        L = expand_param(L, len(sublayers))
        f_ind = self.get_unit('f', *item)
        c_ind = self.get_sublayers('c', sublayers)
        if connect == 'fc':
            mat = self.w_fc_exp
        elif connect == 'cf':
            mat = self.w_cf_exp
        else:
            raise ValueError(f'Invalid connection: {connect}')
        for i, s_ind in enumerate(c_ind):
            s_slice = slice(*tuple(s_ind))
            mat[f_ind, s_slice] += self.c[s_slice] * L[i]

    def study(self, segment, item_list, sublayers, B, Lfc, Lcf):
        """
        Study a list of items.

        Parameters
        ----------
        segment : tuple of str, str
            Sublayer and segment of items to present.

        item_list : numpy.array
            Item indices relative to the segment.

        sublayers : str or list of str
            Sublayers of context to update.

        B : float or numpy.array
            Context updating rate. If an array, specifies a rate for
            each individual study trial.

        Lfc : float or numpy.array
            Learning rate for item to context associations. If an
            array, specifies a learning rate for each individual trial.

        Lcf : float or numpy.array
            Learning rate for context to item associations.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        param = prepare_study_param(item_list.shape[0], len(sublayers), B, Lfc, Lcf)

        # get unit indices
        f_ind = self.get_segment('f', *segment)
        item_ind = (f_ind[0] + item_list).astype(np.dtype('i'))
        c_ind = self.get_sublayers('c', sublayers)

        operations.study(
            self.w_fc_exp,
            self.w_fc_pre,
            self.w_cf_exp,
            self.c,
            self.c_in,
            self.f,
            item_ind,
            c_ind,
            param['B'],
            param['Lfc'],
            param['Lcf'],
        )

    def study_distract(
        self,
        segment,
        item_list,
        sublayers,
        B,
        Lfc,
        Lcf,
        distract_segment,
        distract_list,
        distract_B,
        retention_B,
    ):
        """
        Study a list of items.

        Parameters
        ----------
        segment : tuple of str, str
            Sublayer and segment representing the items to be presented.

        item_list : numpy.array
            Item indices relative to the segment.

        sublayers : str or list of str
            Sublayers to update during study.

        B : float or numpy.array
            Context updating rate. If an array, specifies a rate for
            each individual study trial.

        Lfc : float or numpy.array
            Learning rate for item to context associations. If an
            array, specifies a learning rate for each individual trial.

        Lcf : float or numpy.array
            Learning rate for context to item associations.

        distract_segment : str, str
            Sublayer and segment representing distraction trials.

        distract_list : numpy.array
            Distraction item indices relative to the segment.

        distract_B : float or numpy.array
            Context updating rate for each distraction event before
            each study event. If an array, must be of length
            n_items + 1. Distraction will not be presented on trials i
            where distract_B[i] is zero.
        
        retention_B : float
            Context updating rate for distraction after the last study
            event.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        param = prepare_study_param(
            item_list.shape[0], len(sublayers), B, Lfc, Lcf, distract_B, retention_B
        )

        # get unit indices
        f_ind = self.get_segment('f', *segment)
        item_ind = (f_ind[0] + item_list).astype(np.dtype('i'))
        f_ind_distract = self.get_segment('f', *distract_segment)
        distract_ind = (f_ind_distract[0] + distract_list).astype(np.dtype('i'))
        c_ind = self.get_sublayers('c', sublayers)

        operations.study_distract(
            self.w_fc_exp,
            self.w_fc_pre,
            self.w_cf_exp,
            self.c,
            self.c_in,
            self.f,
            item_ind,
            c_ind,
            param['B'],
            param['Lfc'],
            param['Lcf'],
            distract_ind,
            param['distract_B'],
        )

    def record_study(
        self, 
        segment, 
        item_list, 
        sublayers, 
        B, 
        Lfc, 
        Lcf, 
        distract_segment=None, 
        distract_list=None, 
        distract_B=0, 
        retention_B=0, 
        include=None, 
        exclude=None,
    ):
        """
        Study a list of items and record network states.

        Parameters
        ----------
        segment : tuple of str, str
            Sublayer and segment of items to present.

        item_list : numpy.array
            Item indices relative to the segment.

        sublayers : str or list of str
            Sublayers of context to update.

        B : float or numpy.array
            Context updating rate. If an array, specifies a rate for
            each individual study trial.

        Lfc : float or numpy.array
            Learning rate for item to context associations. If an
            array, specifies a learning rate for each individual trial.

        Lcf : float or numpy.array
            Learning rate for context to item associations.

        distract_segment : str, str
            Sublayer and segment representing distraction trials.

        distract_list : numpy.array
            Distraction item indices relative to the segment.

        distract_B : float or numpy.array
            Context updating rate for each distraction event before
            each study event. If an array, must be of length
            n_items + 1. Distraction will not be presented on trials i
            where distract_B[i] is zero.
        
        retention_B : float
            Context updating rate for distraction after the last study
            event.

        include : list of str, optional
            Network attributes to include in the recorded states.
            Default is to include all attributes.

        exclude : list of str, optional
            Network attributes to exclude from the recorded states.

        Returns
        -------
        state : list of cymr.network.Network
            Copy of the network state after presentation of each item.
        """
        if distract_segment is not None and distract_list is not None:
            distraction = True
        else:
            distraction = False
        
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        param = prepare_study_param(
            item_list.shape[0], len(sublayers), B, Lfc, Lcf, distract_B, retention_B
        )
        state = []
        for i in range(len(item_list)):
            if distraction:
                item = (*distract_segment, distract_list[i])
                self.integrate(item, sublayers, param['distract_B'][i])

            item = (*segment, item_list[i])
            self.present(
                item, sublayers, param['B'][i], param['Lfc'][i], param['Lcf'][i]
            )
            state.append(self.copy(include=include, exclude=exclude))
        
        if distraction:
            item = (*distract_segment, distract_list[-1])
            self.integrate(item, sublayers, param['distract_B'][-1])
        return state

    def record_recall(
        self,
        segment,
        recalls,
        sublayers,
        B,
        T,
        amin=0.000001,
        include=None,
        exclude=None,
    ):
        """
        Simulate a recall sequence and record network states.

        Parameters
        ----------
        segment : tuple of str, str
            Sublayer and segment from which items are recalled.

        recalls : numpy.array
            Index of each recalled item relative to the segment.

        sublayers : str or list of str
            Sublayer(s) of context to update.

        B : float or numpy.array
            Context updating rate. If an array, specifies a rate for
            each individual recall trial.

        T : float
            Decision parameter for choice rule.

        amin : float, optional
            Minimum activation for each not-yet-recalled item on each
            recall attempt.

        include : list of str, optional
            Network attributes to include in the recorded states.
            Default is to include all attributes.

        exclude : list of str, optional
            Network attributes to exclude from the recorded states.

        Returns
        -------
        state : list of cymr.network.Network
            Copy of the network state after each recall attempt.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        recalls = np.array(recalls, dtype=np.dtype('i'))
        n_item = recalls.shape[0]
        n_sub = len(sublayers)
        param = prepare_recall_param(n_item, n_sub, B, T, amin)

        f_ind = self.get_segment('f', *segment)
        exclude_items = np.zeros(self.n_f, dtype=np.dtype('i'))
        state = []
        for i, recall in enumerate(recalls):
            # calculate item support
            operations.cue_item(
                f_ind[0],
                n_item,
                self.w_cf_pre,
                self.w_cf_exp,
                self.w_ff_pre,
                self.w_ff_exp,
                self.f_in,
                self.c,
                exclude_items,
                recalls,
                i,
            )
            operations.apply_softmax(
                f_ind[0], n_item, self.f_in, exclude_items, amin, param['T']
            )

            # update the item layer
            ind = f_ind[0] + recall
            self.f[:] = 0
            self.f[ind] = 1
            exclude_items[ind] = 1

            # save context cue and the item it cued
            state.append(self.copy(include=include, exclude=exclude))

            # update context
            item = (*segment, recall)
            self.integrate(item, sublayers, param['B'][i])
        return state

    def p_recall(self, segment, recalls, sublayers, B, T, p_stop, amin=0.000001):
        """
        Calculate the probability of a specific recall sequence.

        Parameters
        ----------
        segment : tuple of str, str
            Sublayer and segment from which items are recalled.

        recalls : numpy.array
            Index of each recalled item relative to the segment.

        sublayers : str or list of str
            Sublayer(s) of context to update.

        B : float
            Context updating rate after each recalled item.

        T : float
            Decision parameter for choice rule.

        p_stop : numpy.array
            Probability of stopping recall at each output position.

        amin : float, optional
            Minimum activation for each not-yet-recalled item on each
            recall attempt.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        recalls = np.array(recalls, dtype=np.dtype('i'))
        param = prepare_recall_param(recalls.shape[0], len(sublayers), B, T, amin)

        f_ind = self.get_segment('f', *segment)
        start = f_ind[0]
        n_f = f_ind[1] - f_ind[0]
        c_ind = self.get_sublayers('c', sublayers)

        exclude = np.zeros(n_f, dtype=np.dtype('i'))
        p = np.zeros(len(recalls) + 1)
        operations.p_recall(
            start,
            n_f,
            recalls,
            self.w_fc_exp,
            self.w_fc_pre,
            self.w_cf_exp,
            self.w_cf_pre,
            self.w_ff_exp,
            self.w_ff_pre,
            self.f,
            self.f_in,
            self.c,
            self.c_in,
            c_ind,
            exclude,
            amin,
            param['B'],
            param['T'],
            p_stop,
            p,
        )
        return p

    def generate_recall(
        self,
        segment,
        sublayers,
        B,
        T,
        p_stop,
        amin=0.000001,
        filter_recalls=False,
        A1=None,
        A2=None,
    ):
        """
        Generate a sequence of simulated free recall events.

        Parameters
        ----------
        segment : tuple of str, str
            Sublayer and segment to retrieve items from.

        sublayers : str or list of str
            Sublayer(s) of context to update.

        B : float or numpy.ndarray
            Context updating rate after each recall.

        T : float
            Decision parameter for choice rule.

        p_stop : numpy.array
            Probability of stopping at each output position.

        amin : float, optional
            Minimum activation of each not-yet-recalled item on each
            recall attempt.

        filter_recalls : bool, optional
            If true, potential recalls will be filtered based on match
            to context.

        A1 : float, optional
            Intercept mapping context match to an expit to determine
            recovery probability.

        A2 : float, optional
            Slope mapping match to an expit.

        Returns
        -------
        recalls : list of int
            Indices of items recalled at each output position.
        """
        if not isinstance(sublayers, list):
            sublayers = [sublayers]

        # weights to use for recall (assume fixed during recall)
        rec_ind = self.get_segment('f', *segment)
        n_item = rec_ind[1] - rec_ind[0]
        n_sub = len(sublayers)
        param = prepare_recall_param(n_item, n_sub, B, T, amin)

        recalls = []
        exclude = np.zeros(n_item, dtype=np.dtype('i'))
        item_ind = np.arange(n_item)
        for i in range(n_item):
            # stop recall with some probability
            if np.random.rand() < p_stop[i]:
                break

            # calculate item support
            operations.cue_item(
                rec_ind[0],
                n_item,
                self.w_cf_pre,
                self.w_cf_exp,
                self.w_ff_pre,
                self.w_ff_exp,
                self.f_in,
                self.c,
                exclude,
                np.asarray(recalls, dtype=np.dtype('i')),
                i,
            )
            operations.apply_softmax(
                rec_ind[0], n_item, self.f_in, exclude, amin, param['T']
            )

            # select item for recall proportionate to support
            support = self.f_in[rec_ind[0] : rec_ind[1]]
            p_recall = support / np.sum(support)

            if filter_recalls:
                # acceptance probability based on context match
                operations.item_match(
                    rec_ind[0], n_item, self.w_fc_pre, self.w_fc_exp, self.c, self.match
                )
                operations.apply_expit(rec_ind[0], n_item, self.match, A1, A2)

                # recall probabiity is selection + acceptance
                p_recall = p_recall * self.match[rec_ind[0] : rec_ind[1]]

                # rescale probability to vary between 0 and 1
                p_recall = p_recall / np.sum(p_recall)

            if np.any(np.isnan(p_recall)):
                n = np.count_nonzero(exclude == 0)
                p_recall[exclude == 0] = 1 / n
                p_recall[exclude == 1] = 0
                recall = np.random.choice(item_ind, p=p_recall)
            else:
                recall = np.random.choice(item_ind, p=p_recall)
            recalls.append(recall)
            exclude[recall] = 1

            # integrate context associated with the item into context
            item = (*segment, recall)
            self.integrate(item, sublayers, param['B'][i])
        return recalls

    def generate_recall_lba(self, segment, sublayers, time_limit, B, A, b, s, tau):
        """Generate timed free recall using the LBA model."""
        if not isinstance(sublayers, list):
            sublayers = [sublayers]
        rec_ind = self.get_segment('f', *segment)
        n_item = rec_ind[1] - rec_ind[0]
        n_sub = len(sublayers)
        B = expand_param(B, (n_item, n_sub))

        recalls = []
        times = []
        exclude = np.zeros(n_item, dtype=np.dtype('i'))
        t = 0
        for i in range(n_item):
            if t >= time_limit:
                break

            # calculate item support
            operations.cue_item(
                rec_ind[0],
                n_item,
                self.w_cf_pre,
                self.w_cf_exp,
                self.w_ff_pre,
                self.w_ff_exp,
                self.f_in,
                self.c,
                exclude,
                np.asarray(recalls, dtype=np.dtype('i')),
                i,
            )
            support = self.f_in[rec_ind[0] : rec_ind[1]]

            # simulate response competition
            recall, irt = sample_response_lba(A, b, support, s, tau)
            t += irt

            if t <= time_limit:
                # if response happened in time, count it
                recalls.append(recall)
                times.append(t)
                exclude[recall] = 1

                # integrate context associated with the item into context
                item = (*segment, recall)
                self.integrate(item, sublayers, B[i])
        return recalls, times

    def plot(self, ax=None):
        """Plot the current state of the network."""
        if ax is None:
            _, ax = init_plot()
        for name, h in ax.items():
            if not hasattr(self, name):
                continue
            mat = getattr(self, name)
            if mat.ndim == 1:
                mat = mat[None, :]
            h.matshow(mat)
            h.set_aspect('auto')
            h.set_axis_off()
            h.set_title(name, fontsize=12)
