# Standard modules
import copy

import numpy as np
import torch
from scipy.stats import truncnorm

import neuralplayground.agents.whittington_2020_extras.whittington_2020_utils as utils


class Model(torch.nn.Module):
    """
    Model class for TEM model. Inherits from torch.nn.Module, so it can be used as a torch module, and all parameters are
    automatically registered as trainable parameters.
    """

    def __init__(self, params):
        """
        Initialise model with parameters.
        Parameters
        -----------
            params: dict of parameters, usually generated from parameters() in parameters.py
        """
        # First call super class init function to set up torch.nn.Module style model and inherit it's functionality
        super(Model, self).__init__()
        # Copy hyperparameters (e.g. network sizes) from parameter dict, usually generated from parameters() in parameters.py
        self.hyper = copy.deepcopy(params)
        # Create trainable parameters
        self.init_trainable()

    def forward(self, walk, prev_iter=None, prev_M=None):
        """
        Forward pass of TEM model. This consists of a transition, followed by an inference and a generative step, for
        each step of the walk.
        Parameters
        -----------
            walk: list of [place, observation, action] tuples, where place is a list of locations, observation is a one-hot
            vector of sensory observation, and action is a one-hot vector of action taken
            prev_iter: list of Iteration objects, which contain all variables of the model for each step of the previous
            walk. If None, all variables are initialised as zero.
            prev_M: list of memory connectivity matrices for each frequency module. If None, all connectivity matrices are
            initialised as zero.
        Returns
        -----------
            steps: list of Iteration objects, which contain all variables of the model for each step of the walk.
        """
        # The previous iteration may contain walks without action. These are new walks, for which some parameters need to
        # be reset.
        steps = self.init_walks(prev_iter)
        # Forward pass: perform a TEM iteration for each set of [place, observation, action], and produce inferred and
        # generated variables for each step.
        for g, x, a in walk:
            # If there is no previous iteration at all: all walks are new, initialise a whole new iteration object
            if steps is None:
                # Use an Iteration object to set initial values before any real iterations, initialising M, x_inf as zero.
                # Set actions to None blank to indicate there was no previous action
                steps = [self.init_iteration(g, x, [None for _ in range(len(a))], prev_M)]
            # Perform TEM iteration using transition from previous iteration
            L, M, g_gen, p_gen, x_gen, x_logits, x_inf, g_inf, p_inf = self.iteration(
                x, g, steps[-1].a, steps[-1].M, steps[-1].x_inf, steps[-1].g_inf
            )
            # Store this iteration in iteration object in steps list
            steps.append(Iteration(g, x, a, L, M, g_gen, p_gen, x_gen, x_logits, x_inf, g_inf, p_inf))
            # The first step is either a step from a previous walk or initialisiation rubbish, so remove it
        steps = steps[1:]

        # Return steps, which is a list of Iteration objects
        return steps

    def iteration(self, x, locations, a_prev, M_prev, x_prev, g_prev):
        """
        Perform a single iteration of the TEM model. This consists of a transition step, followed by an inference step
        and a generative step.
        Parameters
        -----------
            x: sensory observation, one-hot vector
            locations: list of locations
            a_prev: previous action, one-hot vector
            M_prev: previous memory connectivity matrix
            x_prev: previous sensory experience
            g_prev: previous abstract location
        Returns
        -----------
            L: list of losses for this iteration
            M: list of memory connectivity matrices for this iteration
            gt_gen: list of generated abstract (grid cell) locations for this iteration
            p_gen: list of generated place cell locations for this iteration
            x_gen: list of generated observations for this iteration
            x_logits: list of logits for generated observations for this iteration
            x_inf: list of inferred observations for this iteration
            g_inf: list of inferred abstract (grid cell) locations for this iteration
            p_inf: list of inferred grounded (place cell) locations for this iteration
        """
        # First, do the transition step, as it will be necessary for both the inference and generative part of the model
        gt_gen, gt_inf = self.gen_g(a_prev, g_prev, locations)
        # Run inference model: infer grounded location p_inf (hippocampus), abstract location g_inf (entorhinal).
        # Also keep filtered sensory observation (x_inf), and retrieved grounded location p_inf_x
        x_inf, g_inf, p_inf_x, p_inf = self.inference(x, locations, M_prev, x_prev, gt_inf)
        # Run generative model: since generative model is only used for training purposes, it will generate from
        # *inferred* variables instead of *generated* variables (as it would when used for generation)
        x_gen, x_logits, p_gen = self.generative(M_prev, p_inf, g_inf, gt_gen)
        # Update generative memory with generated and inferred grounded location.
        M = [self.hebbian(M_prev[0], torch.cat(p_inf, dim=1), torch.cat(p_gen, dim=1))]
        # If using memory for grounded location inference: append inference memory
        if self.hyper[
            "use_p_inf"
        ]:  # Inference memory is identical to generative memory if using common memory, and updated separatedly if not
            M.append(
                M[0]
                if self.hyper["common_memory"]
                else self.hebbian(
                    M_prev[1], torch.cat(p_inf, dim=1), torch.cat(p_inf_x, dim=1), do_hierarchical_connections=False
                )
            )
        # Calculate loss of this step
        L = self.loss(gt_gen, p_gen, x_logits, x, g_inf, p_inf, p_inf_x)
        # Return all iteration values
        return L, M, gt_gen, p_gen, x_gen, x_logits, x_inf, g_inf, p_inf

    def inference(self, x, locations, M_prev, x_prev, g_gen):
        """
        Perform inference step of the TEM model.
        Parameters
        -----------
            x: sensory observation, one-hot vector
            locations: list of locations
            M_prev: previous memory connectivity matrix
            x_prev: previous sensory experience
            g_gen: previous abstract location
        Returns
        -----------
            x_f: filtered sensory experience
            g: inferred abstract (grid cell) locations
            p_x: retrieved grounded (place cell) locations from memory by sensory experience
            p: inferred grounded (place cell) locations
        """
        # Compress sensory observation from one-hot to two-hot (or alternatively, whatever an MLP makes of it)
        x_c = self.f_c(x)
        # Temporally filter sensory observation by mixing it with previous experience
        x_f = self.x_prev2x(x_prev, x_c)
        # Prepare sensory experience for input to memory by normalisation and weighting
        x_ = self.x2x_(x_f)
        # Retrieve grounded location from memory by doing pattern completion on current sensory experience
        p_x = (
            self.attractor(x_, M_prev[1], retrieve_it_mask=self.hyper["p_retrieve_mask_inf"])
            if self.hyper["use_p_inf"]
            else None
        )
        # Infer abstract location by combining previous abstract location and grounded location retrieved from
        # memory by current sensory experience
        g = self.inf_g(p_x, g_gen, x, locations)
        # Prepare abstract location for input to memory by downsampling and weighting
        g_ = self.g2g_(g)
        # Infer grounded location from sensory experience and inferred abstract location
        p = self.inf_p(x_, g_)
        # Return variables in order that they were created
        return x_f, g, p_x, p

    def generative(self, M_prev, p_inf, g_inf, g_gen):
        """
        Perform generative step of the TEM model.
        Parameters
        -----------
            M_prev: previous memory connectivity matrix
            p_inf: inferred grounded (place cell) locations
            g_inf: inferred abstract (grid cell) locations
            g_gen: previous abstract locations
        Returns
        -----------
            x_p: generated observation from inferred grounded location
            x_g: generated observation from grounded location retrieved from inferred abstract
            location
            x_gt: generated observation from grounded location retrieved from abstract location
            by transitioning
            x_p_logits: logits for generated observation from inferred grounded location
            x_g_logits: logits for generated observation from grounded location retrieved from
            inferred abstract location
            x_gt_logits: logits for generated observation from grounded location retrieved from
            abstract location by transitioning
            p_g_inf: grounded location retrieved from inferred abstract location
        """
        # Generate observation from inferred grounded location, using only the highest frequency.
        # Also keep non-softmaxed logits which are used in the loss later
        x_p, x_p_logits = self.gen_x(p_inf[0])
        # Retrieve grounded location from memory by pattern completion on inferred abstract
        # location
        p_g_inf = self.gen_p(g_inf, M_prev[0])  # was p_mem_gen
        # And generate observation from the grounded location retrieved from inferred abstract
        # location
        x_g, x_g_logits = self.gen_x(p_g_inf[0])
        # Retreive grounded location from memory by pattern completion on abstract location by
        # transitioning
        p_g_gen = self.gen_p(g_gen, M_prev[0])
        # Generate observation from sampled grounded location
        x_gt, x_gt_logits = self.gen_x(p_g_gen[0])
        # Return all generated observations and their corresponding logits
        return (x_p, x_g, x_gt), (x_p_logits, x_g_logits, x_gt_logits), p_g_inf

    def loss(self, g_gen, p_gen, x_logits, x, g_inf, p_inf, p_inf_x):
        """
        Calculate losses of the current TEM iteration for each of the generated and inferred
        variable.
        Parameters
        -----------
            g_gen: previous abstract (grid cell) locations
            p_gen: previous grounded (place cell) locations
            x_logits: logits for generated observation from grounded location retrieved from abstract
            location by transitioning
            x: sensory observation, one-hot vector
            g_inf: inferred abstract (grid cell) locations
            p_inf: inferred grounded (place cell) locations
            p_inf_x: retrieved grounded (place cell) locations from memory by sensory experience
        Returns
        -----------
            L: list of losses for this iteration
        """
        # Calculate loss function, separately for each component because you might want to reweight
        # contributions later
        # L_p_gen is squared error loss between inferred grounded location and grounded location retrieved
        # from inferred abstract location
        L_p_g = torch.sum(torch.stack(utils.squared_error(p_inf, p_gen), dim=0), dim=0)
        # L_p_inf is squared error loss between inferred grounded location and grounded location retrieved
        # from sensory experience
        L_p_x = (
            torch.sum(torch.stack(utils.squared_error(p_inf, p_inf_x), dim=0), dim=0)
            if self.hyper["use_p_inf"]
            else torch.zeros_like(L_p_g)
        )
        # L_g is squared error loss between generated abstract location and inferred abstract location
        L_g = torch.sum(torch.stack(utils.squared_error(g_inf, g_gen), dim=0), dim=0)
        # L_x is a cross-entropy loss between sensory experience and different model predictions. First get
        # true labels from sensory experience
        labels = torch.argmax(x, 1)
        # L_x_gen: losses generated by generative model from g_prev -> g -> p -> x
        L_x_gen = utils.cross_entropy(x_logits[2], labels)
        # L_x_g: Losses generated by generative model from g_inf -> p -> x
        L_x_g = utils.cross_entropy(x_logits[1], labels)
        # L_x_p: Losses generated by generative model from p_inf -> x
        L_x_p = utils.cross_entropy(x_logits[0], labels)
        # L_reg are regularisation losses, L_reg_g on L2 norm of g
        L_reg_g = torch.sum(torch.stack([torch.sum(g**2, dim=1) for g in g_inf], dim=0), dim=0)
        # And L_reg_p regularisation on L1 norm of p
        L_reg_p = torch.sum(torch.stack([torch.sum(torch.abs(p), dim=1) for p in p_inf], dim=0), dim=0)
        # Return total loss as list of losses, so you can possibly reweight them
        L = [L_p_g, L_p_x, L_x_gen, L_x_g, L_x_p, L_g, L_reg_g, L_reg_p]
        return L

    def init_trainable(self):
        """
        Initialise all trainable parameters of the TEM model. This is done in a separate function so that
        it can be called again after loading a model from file.
        """
        # Scale factor in Laplacian transform for each frequency module. High frequency comes first, low
        # frequency comes last. Learn inverse sigmoid instead of scale factor directly, so domain of alpha is -inf, inf
        self.alpha = torch.nn.ParameterList(
            [
                torch.nn.Parameter(
                    torch.tensor(np.log(self.hyper["f_initial"][f] / (1 - self.hyper["f_initial"][f])), dtype=torch.float)
                )
                for f in range(self.hyper["n_f"])
            ]
        )
        # Entorhinal preference weights
        self.w_x = torch.nn.Parameter(torch.tensor(1.0))
        # Entorhinal preference bias
        self.b_x = torch.nn.Parameter(torch.zeros(self.hyper["n_x_c"]))
        # Frequency module specific scaling of sensory experience before input to hippocampus
        self.w_p = torch.nn.ParameterList([torch.nn.Parameter(torch.tensor(1.0)) for f in range(self.hyper["n_f"])])
        # Initial activity of abstract location cells when entering a new environment, like a prior on g.
        # Initialise with truncated normal
        self.g_init = torch.nn.ParameterList(
            [
                torch.nn.Parameter(
                    torch.tensor(
                        truncnorm.rvs(-2, 2, size=self.hyper["n_g"][f], loc=0, scale=self.hyper["g_init_std"]),
                        dtype=torch.float,
                    )
                )
                for f in range(self.hyper["n_f"])
            ]
        )
        # Log of standard deviation of abstract location cells when entering a new environment; standard
        # deviation of the prior on g. Initialise with truncated normal
        self.logsig_g_init = torch.nn.ParameterList(
            [
                torch.nn.Parameter(
                    torch.tensor(
                        truncnorm.rvs(-2, 2, size=self.hyper["n_g"][f], loc=0, scale=self.hyper["g_init_std"]),
                        dtype=torch.float,
                    )
                )
                for f in range(self.hyper["n_f"])
            ]
        )
        # MLP for transition weights (not in paper, but recommended by James so you can learn about
        # similarities between actions). Size is given by grid connections
        self.MLP_D_a = MLP(
            [self.hyper["n_actions"] for _ in range(self.hyper["n_f"])],
            [
                sum(
                    [
                        self.hyper["n_g"][f_from]
                        for f_from in range(self.hyper["n_f"])
                        if self.hyper["g_connections"][f_to][f_from]
                    ]
                )
                * self.hyper["n_g"][f_to]
                for f_to in range(self.hyper["n_f"])
            ],
            activation=[torch.tanh, None],
            hidden_dim=[self.hyper["d_hidden_dim"] for _ in range(self.hyper["n_f"])],
            bias=[True, False],
        )
        # Initialise the hidden to output weights as zero, so initially you simply keep the current abstract
        # location to predict the next abstract location
        self.MLP_D_a.set_weights(1, 0.0)
        # Transition weights without specifying an action for use in generative model with shiny objects
        self.D_no_a = torch.nn.ParameterList(
            [
                torch.nn.Parameter(
                    torch.zeros(
                        sum(
                            [
                                self.hyper["n_g"][f_from]
                                for f_from in range(self.hyper["n_f"])
                                if self.hyper["g_connections"][f_to][f_from]
                            ]
                        )
                        * self.hyper["n_g"][f_to]
                    )
                )
                for f_to in range(self.hyper["n_f"])
            ]
        )
        # MLP for standard deviation of transition sample
        self.MLP_sigma_g_path = MLP(
            self.hyper["n_g"],
            self.hyper["n_g"],
            activation=[torch.tanh, torch.exp],
            hidden_dim=[2 * g for g in self.hyper["n_g"]],
        )
        # MLP for standard devation of grounded location from retrieved memory sample
        self.MLP_sigma_p = MLP(self.hyper["n_p"], self.hyper["n_p"], activation=[torch.tanh, torch.exp])
        # MLP to generate mean of abstract location from downsampled abstract location, obtained by summing
        # grounded location over sensory preferences in inference model
        self.MLP_mu_g_mem = MLP(self.hyper["n_g_subsampled"], self.hyper["n_g"], hidden_dim=[2 * g for g in self.hyper["n_g"]])
        # Initialise weights in last layer of MLP_mu_g_mem as truncated normal for each frequency module
        self.MLP_mu_g_mem.set_weights(
            -1,
            [
                torch.tensor(
                    truncnorm.rvs(
                        -2, 2, size=list(self.MLP_mu_g_mem.w[f][-1].weight.shape), loc=0, scale=self.hyper["g_mem_std"]
                    ),
                    dtype=torch.float,
                )
                for f in range(self.hyper["n_f"])
            ],
        )
        # MLP to generate standard deviation of abstract location from two measures (generated observation
        # error and inferred abstract location vector norm) of memory quality
        self.MLP_sigma_g_mem = MLP(
            [2 for _ in self.hyper["n_g_subsampled"]],
            self.hyper["n_g"],
            activation=[torch.tanh, torch.exp],
            hidden_dim=[2 * g for g in self.hyper["n_g"]],
        )
        # MLP to generate mean of abstract location directly from shiny object presence. Outputs to object
        # vector cell modules if they're separated, else to all abstract location modules
        self.MLP_mu_g_shiny = MLP(
            [1 for _ in range(self.hyper["n_f_ovc"] if self.hyper["separate_ovc"] else self.hyper["n_f"])],
            [n_g for n_g in self.hyper["n_g"][(self.hyper["n_f_g"] if self.hyper["separate_ovc"] else 0) :]],
            hidden_dim=[2 * n_g for n_g in self.hyper["n_g"][(self.hyper["n_f_g"] if self.hyper["separate_ovc"] else 0) :]],
        )
        # MLP to generate standard deviation of abstract location directly from shiny object presence.
        # Outputs to object vector cell modules if they're separated, else to all abstract location modules
        self.MLP_sigma_g_shiny = MLP(
            [1 for _ in range(self.hyper["n_f_ovc"] if self.hyper["separate_ovc"] else self.hyper["n_f"])],
            [n_g for n_g in self.hyper["n_g"][(self.hyper["n_f_g"] if self.hyper["separate_ovc"] else 0) :]],
            hidden_dim=[2 * n_g for n_g in self.hyper["n_g"][(self.hyper["n_f_g"] if self.hyper["separate_ovc"] else 0) :]],
            activation=[torch.tanh, torch.exp],
        )
        # MLP for decompressing highest frequency sensory experience to sensory observation
        self.MLP_c_star = MLP(self.hyper["n_x_f"][0], self.hyper["n_x"], hidden_dim=20 * self.hyper["n_x_c"])

    def init_iteration(self, g, x, a, M):
        """
        Initialise a new iteration of the TEM model.
        Parameters
        -----------
            g: previous abstract (grid cell) locations
            x: sensory observation, one-hot vector
            a: previous action, one-hot vector
            M: previous memory connectivity matrix
        Returns
        -----------
            iteration: Iteration object, which contains all variables of the model for this iteration.
        """
        # On the very first iteration, update the batch size based on the data. This is useful when
        # doing analysis on the network with different batch sizes compared to training
        self.hyper["batch_size"] = x.shape[0]
        # Initalise hebbian memory connectivity matrix [M_gen, M_inf] if it wasn't initialised yet
        if M is None:
            # Create new empty memory dict for generative network: zero connectivity matrix M_0, then empty
            # list of the memory vectors a and b for each iteration for efficient hebbian memory computation
            M = [torch.zeros((self.hyper["batch_size"], sum(self.hyper["n_p"]), sum(self.hyper["n_p"])), dtype=torch.float)]
            # Append inference memory only if memory is used in grounded location inference
            if self.hyper["use_p_inf"]:
                # If inference and generative network share common memory: reuse same connectivity, and
                # same memory vectors. Else, create a new empty memory list for inference network
                M.append(
                    M[0]
                    if self.hyper["common_memory"]
                    else torch.zeros(
                        (self.hyper["batch_size"], sum(self.hyper["n_p"]), sum(self.hyper["n_p"])), dtype=torch.float
                    )
                )
                # Initialise previous abstract location by stacking abstract location prior
        g_inf = [torch.stack([self.g_init[f] for _ in range(self.hyper["batch_size"])]) for f in range(self.hyper["n_f"])]
        # Initialise previous sensory experience with zeros, as there is no data yet for
        # temporal smoothing
        x_inf = [torch.zeros((self.hyper["batch_size"], self.hyper["n_x_f"][f])) for f in range(self.hyper["n_f"])]
        # And construct new iteration for that g, x, a, and M
        return Iteration(g=g, x=x, a=a, M=M, x_inf=x_inf, g_inf=g_inf)

    def init_walks(self, prev_iter):
        """
        Initialise a new walk of the TEM model.
        Parameters
        -----------
            prev_iter: list of Iteration objects, which contain all variables of the model for each step of the
            previous walk. If None, all variables are initialised as zero.
        Returns
        -----------
            prev_iter: list of Iteration objects, which contain all variables of the model for each step of the
            previous walk. If None, all variables are initialised as zero.
        """
        # Only reset parameters for previous iteration if a previous iteration was actually provided - if
        # it wasn't, all parameters will be reset when creating a fresh Iteration object in init_iteration
        if prev_iter is not None:
            # The supplied previous iteration might have new walks starting, with empty actions. For these
            # walks some parameters need to be reset
            for a_i, a in enumerate(prev_iter[0].a):
                # A new walk is indicated by having a None action in the previous iteration
                if a is None:
                    # Reset the initial connectivity matrix for this walk
                    for M in prev_iter[0].M:
                        M[a_i, :, :] = 0
                        # Reset the abstract location for this walk
                    for f, g_inf in enumerate(prev_iter[0].g_inf):
                        g_inf[a_i, :] = self.g_init[f]
                    # Reset the sensory experience for this walk
                    for f, x_inf in enumerate(prev_iter[0].x_inf):
                        x_inf[a_i, :] = torch.zeros(self.hyper["n_x_f"][f])
        # Return the iteration with reset parameters (or simply the empty array if prev_iter was empty)
        return prev_iter

    def gen_g(self, a_prev, g_prev, locations):
        """
        Perform transition step of the TEM model. This consists of a transition from previous abstract location
        to new abstract location using weights specific to
        action taken for each frequency module.
        Parameters
        -----------
            a_prev: previous action, one-hot vector
            g_prev: previous abstract (grid cell) locations
            locations: list of locations
        Returns
        -----------
            g_gen: generated abstract (grid cell) locations
            g: abstract (grid cell) locations
        """
        # Transition from previous abstract location to new abstract location using weights specific to action
        # taken for each frequency module
        mu_g = self.f_mu_g_path(a_prev, g_prev)
        sigma_g = self.f_sigma_g_path(a_prev, g_prev)
        # Either sample new abstract location g or simply take the mean of distribution in noiseless case.
        g = [
            mu_g[f] + sigma_g[f] * np.random.randn() if self.hyper["do_sample"] else mu_g[f] for f in range(self.hyper["n_f"])
        ]
        # But for environments with shiny objects, the transition to the new abstract location shouldn't have
        #  access to the action direction in the generative model
        shiny_envs = [location["shiny"] is not None for location in locations]
        # If there are any shiny environments, the abstract locations for the generative model will need to be
        # re-calculated without providing actions for those
        g_gen = self.f_mu_g_path(a_prev, g_prev, no_direc=shiny_envs) if any(shiny_envs) else g
        # Return generated abstract location after transition
        return g_gen, (g, sigma_g)

    def gen_p(self, g, M_prev):
        """
        Generate grounded location from abstract location by pattern completion on abstract location. This is
        used in the generative model to generate observations from grounded location.
        Parameters
        -----------
            g: abstract (grid cell) locations
            M_prev: previous memory connectivity matrix
        Returns
        -----------
            p: generated grounded (place cell) locations
        """
        # We want to use g as an index for memory retrieval, but it doesn't have the right dimensions (these
        # are grid cells, we need place cells). We need g_ instead
        g_ = self.g2g_(g)
        # Retreive memory: do pattern completion on abstract location to get grounded location
        mu_p = self.attractor(g_, M_prev, retrieve_it_mask=self.hyper["p_retrieve_mask_gen"])
        sigma_p = self.f_sigma_p(mu_p)
        # Either sample new grounded location p or simply take the mean of distribution in noiseless case
        p = [
            mu_p[f] + sigma_p[f] * np.random.randn() if self.hyper["do_sample"] else mu_p[f] for f in range(self.hyper["n_f"])
        ]
        # Return pattern-completed grounded location p after memory retrieval
        return p

    def gen_x(self, p):
        """
        Generate observation from grounded location by MLP. This is used in the generative model to generate
        observations from grounded location.
        Parameters
        -----------
            p: grounded (place cell) locations
        Returns
        -----------
            x: generated observation
            logits: logits for generated observation
        """
        # Get categorical distribution over observations from grounded location
        # If you actually want to sample observation, you need a reparaterisation trick for categorical distributions
        # Sampling would be the correct way to do this, since observations are discrete, and it's also what the
        # TEM paper says
        # However, it looks like you could also get away with using categorical distribution directly as an
        # approximation of the one-hot observations
        if self.hyper["do_sample"]:
            x, logits = self.f_x(p)  # This is a placeholder! Should be done using reparameterisation trick (like
            # https://blog.evjang.com/2016/11/tutorial-categorical-variational.html)
        else:
            x, logits = self.f_x(p)
        # Return one-hot (or almost one-hot...) observation obtained from grounded location, and also
        # the non-softmaxed logits
        return x, logits

    def inf_g(self, p_x, g_gen, x, locations):
        """
        Infer abstract location from the combination of [grounded location retrieved from memory by sensory
        experience] and [previous abstract location and action (path integration)]
        Parameters
        -----------
            p_x: retrieved grounded (place cell) locations from memory by sensory experience
            g_gen: previous abstract (grid cell) locations
            x: sensory observation, one-hot vector
            locations: list of locations
        Returns
        -----------
            mu_g: inferred abstract (grid cell) locations
            sigma_g: standard deviation of inferred abstract (grid cell) locations
        """
        # Infer abstract location from the combination of [grounded location retrieved from memory by sensory
        # experience] ...
        if self.hyper["use_p_inf"]:
            # Not in paper, but makes sense from symmetry with f_x: first get g from p by "summing over sensory
            # preferences" g = p * W_repeat^T
            g_downsampled = [torch.matmul(p_x[f], torch.t(self.hyper["W_repeat"][f])) for f in range(self.hyper["n_f"])]
            # Then use abstract location after summing over sensory preferences as input to MLP to obtain the
            # inferred abstract location from memory
            mu_g_mem = self.f_mu_g_mem(g_downsampled)
            # Not in paper, but this greatly improves zero-shot inference: provide the uncertainty function of the
            # inferred abstract location with measures of memory quality
            with torch.no_grad():
                # For the first measure, use the grounded location inferred from memory to generate an observation
                x_hat, x_hat_logits = self.gen_x(p_x[0])
                # Then calculate the error between the generated observation and the actual observation: if the memory
                # is working well, this error should be small
                err = utils.squared_error(x, x_hat)
            # The second measure is the vector norm of the inferred abstract location; good memories should have similar
            # vector norms. Concatenate the two measures as input for the abstract location uncertainty function
            sigma_g_input = [
                torch.cat((torch.sum(g**2, dim=1, keepdim=True), torch.unsqueeze(err, dim=1)), dim=1) for g in mu_g_mem
            ]
            # Not in paper, but recommended by James for stability: get final mean of inferred abstract location by
            # clamping activations between -1 and 1
            mu_g_mem = self.f_g_clamp(mu_g_mem)
            # And get standard deviation/uncertainty of inferred abstract location by providing uncertainty function with
            # memory quality measures
            sigma_g_mem = self.f_sigma_g_mem(sigma_g_input)
            # ... and [previous abstract location and action (path integration)]
        mu_g_path = g_gen[0]
        sigma_g_path = g_gen[1]
        # Infer abstract location by combining previous abstract location and grounded location retrieved from memory by
        # current sensory experience
        mu_g, sigma_g = [], []
        for f in range(self.hyper["n_f"]):
            if self.hyper["use_p_inf"]:
                # Then get full gaussian distribution of inferred abstract location by calculating precision weighted mean
                mu, sigma = utils.inv_var_weight([mu_g_path[f], mu_g_mem[f]], [sigma_g_path[f], sigma_g_mem[f]])
            else:
                # Or simply completely ignore the inference memory here, to test if things are working
                mu, sigma = mu_g_path[f], sigma_g_path[f]
            # Append mu and sigma to list for all frequency modules
            mu_g.append(mu)
            sigma_g.append(sigma)
        # Finally (though not in paper), also add object vector cell information to inferred abstract location for
        # environments with shiny objects
        shiny_envs = [location["shiny"] is not None for location in locations]
        if any(shiny_envs):
            # Find for which environments the current location has a shiny object
            shiny_locations = torch.unsqueeze(
                torch.stack(
                    [
                        torch.tensor(location["shiny"], dtype=torch.float)
                        for location in locations
                        if location["shiny"] is not None
                    ]
                ),
                dim=-1,
            )
            # Get abstract location for environments with shiny objects and feed to each of the object vector cell modules
            mu_g_shiny = self.f_mu_g_shiny(
                [shiny_locations for _ in range(self.hyper["n_f_g"] if self.hyper["separate_ovc"] else self.hyper["n_f"])]
            )
            sigma_g_shiny = self.f_sigma_g_shiny(
                [shiny_locations for _ in range(self.hyper["n_f_g"] if self.hyper["separate_ovc"] else self.hyper["n_f"])]
            )
            # Update only object vector modules with shiny-inferred abstract location: start from offset if object vector
            # modules are separate
            module_start = self.hyper["n_f_g"] if self.hyper["separate_ovc"] else 0
            # Inverse variance weighting is associative, so I can just do additional inverse variance weighting to the
            # previously obtained mu and sigma - but only for object vector cell modules!
            for f in range(module_start, self.hyper["n_f"]):
                # Add inferred abstract location from shiny objects to previously obtained position, only for environments
                # with shiny objects
                mu, sigma = utils.inv_var_weight(
                    [mu_g[f][shiny_envs, :], mu_g_shiny[f - module_start]],
                    [sigma_g[f][shiny_envs, :], sigma_g_shiny[f - module_start]],
                )
                # In order to update only the environments with shiny objects, without in-place value assignment,
                # construct a mask of shiny environments
                mask = torch.zeros_like(mu_g[f], dtype=torch.bool)
                mask[shiny_envs, :] = True
                # Use mask to update the shiny environment entries in inferred abstract locations
                mu_g[f] = mu_g[f].masked_scatter(mask, mu)
                sigma_g[f] = sigma_g[f].masked_scatter(mask, sigma)
                # Either sample inferred abstract location from combined (precision weighted) distribution or just take
                # mean
        g = [
            mu_g[f] + sigma_g[f] * np.random.randn() if self.hyper["do_sample"] else mu_g[f] for f in range(self.hyper["n_f"])
        ]
        # Return abstract location inferred from grounded location from memory and previous abstract location
        return g

    def inf_p(self, x_, g_):
        """
        Infer grounded location from sensory experience and inferred abstract location for each module
        Parameters
        -----------
            x_: sensory input to memory
            g_: abstract (grid cell) locations
        Returns
        -----------
            p: inferred grounded (place cell) locations
        """
        # Infer grounded location from sensory experience and inferred abstract location for each module
        p = []
        # Use the same transformation for each frequency module: leaky relu for sparsity
        for f in range(self.hyper["n_f"]):
            mu_p = self.f_p(g_[f] * x_[f])  # This is element-wise multiplication
            sigma_p = 0  # Unclear from paper (typo?). Some undefined function f that takes two arguments: f(f_n(x),g)
            # Either sample inferred grounded location or just take mean
            if self.hyper["do_sample"]:
                p.append(mu_p + sigma_p * np.random.randn())
            else:
                p.append(mu_p)
        # Return new memory constructed from sensory experience and inferred abstract location
        return p

    def x_prev2x(self, x_prev, x_c):
        """
        Perform temporal filtering on sensory experience for each frequency module
        Parameters
        -----------
            x_prev: previous sensory experience
            x_c: current sensory experience
        Returns
        -----------
            x: filtered sensory experience
        """
        # Calculate factor for filtering from sigmoid of learned parameter
        alpha = [torch.nn.Sigmoid()(self.alpha[f]) for f in range(self.hyper["n_f"])]
        # Do exponential temporal filtering for each frequency modulemod
        x = [(1 - alpha[f]) * x_prev[f] + alpha[f] * x_c for f in range(self.hyper["n_f"])]
        return x

    def x2x_(self, x):
        """
        Prepare sensory input for input to memory by weighting and normalisation for each frequency module
        Parameters
        -----------
            x: sensory observation, one-hot vector
        Returns
        -----------
            x_: weighted and normalised sensory input to memory
        """
        # Prepare sensory input for input to memory by weighting and normalisation for each frequency module
        # Get normalised sensory input for each frequency module
        normalised = self.f_n(x)
        # Then reshape and reweight (use sigmoid to keep weight between 0 and 1) each frequency module separately:
        #  matrix multiplication by W_tile prepares x for outer product with g by element-wise multiplication
        x_ = [
            torch.nn.Sigmoid()(self.w_p[f]) * torch.matmul(normalised[f], self.hyper["W_tile"][f])
            for f in range(self.hyper["n_f"])
        ]
        return x_

    def g2g_(self, g):
        """
        Prepare abstract location for input to memory by reshaping and down-sampling for each frequency module
        Parameters
        -----------
            g: abstract (grid cell) locations
        Returns
        -----------
            g_: downsampled abstract (grid cell) locations
        """
        # Prepares abstract location for input to memory by reshaping and down-sampling for each frequency module
        # Get downsampled abstract location for each frequency module
        downsampled = self.f_g(g)
        # Then reshape and reweight each frequency module separately
        g_ = [torch.matmul(downsampled[f], self.hyper["W_repeat"][f]) for f in range(self.hyper["n_f"])]
        return g_

    def f_mu_g_path(self, a_prev, g_prev, no_direc=None):
        """
        Perform transition step of the TEM model. This consists of a transition from previous abstract location
        to new abstract location using weights specific to
        action taken for each frequency module.
        Parameters
        -----------
            a_prev: previous action, one-hot vector
            g_prev: previous abstract (grid cell) locations
            no_direc: list of booleans indicating for which walks the transition direction should be omitted
            (e.g. no shiny objects, or in inference model: set to all false)
        Returns
        -----------
            g_step: new abstract (grid cell) locations
        """
        # If there are no environments where the transition direction needs to be omitted (e.g. no shiny objects,
        # or in inference model: set to all false
        no_direc = [False for _ in a_prev] if no_direc is None else no_direc
        # Remove all Nones from a_prev: these are walks where there was no previous action, so no step needs to
        # be calculated for those
        a_prev_step = [a if a is not None else 0 for a in a_prev]
        # And also keep track of which walks these valid step actions are for
        a_do_step = [a is not None for a in a_prev]
        # Transform list of actions into batch of one-hot row vectors.
        if self.hyper["has_static_action"]:
            # If this world has static actions: whenever action 0 (standing still) appears, the action vector
            # should be all zeros. All other actions should have a 1 in the label-1 entry
            a = torch.zeros((len(a_prev_step), self.hyper["n_actions"])).scatter_(
                1,
                torch.clamp(torch.tensor(a_prev_step).unsqueeze(1) - 1, min=0),
                1.0 * (torch.tensor(a_prev_step).unsqueeze(1) > 0),
            )
        else:
            # Without static actions: each action label should become a one-hot vector for that label
            a = torch.zeros((len(a_prev_step), self.hyper["n_actions"])).scatter_(
                1, torch.tensor(a_prev_step).unsqueeze(1), 1.0
            )
        # Get vector of transition weights by feeding actions into MLP
        D_a = self.MLP_D_a([a for _ in range(self.hyper["n_f"])])
        # Replace transition weights by non-directional transition weights in environments where transition direction
        # needs to be omitted (can set only if any no_direc)
        for f in range(self.hyper["n_f"]):
            D_a[f][no_direc, :] = self.D_no_a[f]
        # Reshape transition weight vector into transition matrix. The number of rows in the transition matrix is given
        # by the incoming abstract location connections for each frequency module
        D_a = [
            torch.reshape(
                D_a[f_to],
                (
                    -1,
                    sum(
                        [
                            self.hyper["n_g"][f_from]
                            for f_from in range(self.hyper["n_f"])
                            if self.hyper["g_connections"][f_to][f_from]
                        ]
                    ),
                    self.hyper["n_g"][f_to],
                ),
            )
            for f_to in range(self.hyper["n_f"])
        ]
        # Select the frequency modules of the previous abstract location that are connected to each frequency module, to
        g_in = [
            torch.unsqueeze(
                torch.cat(
                    [g_prev[f_from] for f_from in range(self.hyper["n_f"]) if self.hyper["g_connections"][f_to][f_from]], dim=1
                ),
                1,
            )
            for f_to in range(self.hyper["n_f"])
        ]
        # Reshape transition weight vector into transition matrix. The number of rows in the transition matrix is given by the
        # incoming abstract location connections for each frequency module
        delta = [torch.squeeze(torch.matmul(g, T)) for g, T in zip(g_in, D_a)]
        # Not in the paper, but recommended by James for stability: use inferred code as *difference* in abstract location.
        # Calculate new abstract location from previous abstract location and difference
        g_step = [g + d if g.dim() > 1 else torch.unsqueeze(g + d, 0) for g, d in zip(g_prev, delta)]
        # Not in paper, but recommended by James for stability: clamp activations between -1 and 1
        g_step = self.f_g_clamp(g_step)
        # Build new abstract location from result of transition if there was one, or from prior on abstract location
        # if there wasn't
        return [
            torch.stack([g_step[f][batch_i, :] if do_step else self.g_init[f] for batch_i, do_step in enumerate(a_do_step)])
            for f in range(self.hyper["n_f"])
        ]

    def f_sigma_g_path(self, a_prev, g_prev):
        """
        Use multi layer perceptron to generate standard deviation from all previous abstract locations, including those
        that were just initialised and not real previous locations.
        Parameters
        -----------
            a_prev: previous action, one-hot vector
            g_prev: previous abstract (grid cell) locations
        Returns
        -----------
            sigma_g: standard deviation of new abstract (grid cell) locations
        """
        # Keep track of which walks these valid step actions are for
        a_do_step = [a is not None for a in a_prev]
        from_g = self.MLP_sigma_g_path(g_prev)
        # And take exponent to get prior sigma for the walks that didn't have a previous location
        from_prior = [torch.exp(logsig) for logsig in self.logsig_g_init]
        # Now select the standard deviation generated from the previous abstract location if there was one, and the prior
        # standard deviation on abstract location otherwise
        return [
            torch.stack([from_g[f][batch_i, :] if do_step else from_prior[f] for batch_i, do_step in enumerate(a_do_step)])
            for f in range(self.hyper["n_f"])
        ]

    def f_mu_g_mem(self, g_downsampled):
        """
        Use multi layer perceptron to generate mean of abstract location from down-sampled abstract location, obtained by
        summing over sensory dimension of grounded location.
        Parameters
        -----------
            g_downsampled: downsampled abstract (grid cell) locations
        Returns
        -----------
            mu_g_mem: inferred abstract (grid cell) locations
        """
        return self.MLP_mu_g_mem(g_downsampled)

    def f_sigma_g_mem(self, g_downsampled):
        """
        Use multi layer perceptron to generate standard deviation of abstract location from down-sampled abstract location,
        obtained by summing over sensory dimension of grounded location.
        Parameters
        -----------
            g_downsampled: downsampled abstract (grid cell) locations
        Returns
        -----------
            sigma_g_mem: standard deviation of inferred abstract (grid cell) locations
        """
        sigma = self.MLP_sigma_g_mem(g_downsampled)
        # Not in paper, but also offset this sigma over training, so you can reduce influence of inferred p early on
        return [sigma[f] + self.hyper["p2g_scale_offset"] * self.hyper["p2g_sig_val"] for f in range(self.hyper["n_f"])]

    def f_mu_g_shiny(self, shiny):
        """
        Use multi layer perceptron to generate mean of abstract location from boolean location shiny-ness. Outputs to
        object vector cell modules if they're separated,
        else to all abstract location modules
        Parameters
        -----------
            shiny: boolean location shiny-ness
        Returns
        -----------
            mu_g: inferred abstract (grid cell) locations
        """
        mu_g = self.MLP_mu_g_shiny(shiny)
        # Take absolute because James wants object vector cells to be positive
        mu_g = [torch.abs(mu) for mu in mu_g]
        # Then apply clamp and leaky relu to get object vector module activations, like it's done for ground location
        # activations
        g = self.f_p(mu_g)
        return g

    def f_sigma_g_shiny(self, shiny):
        """
        Use multi layer perceptron to generate standard deviation of abstract location from boolean location shiny-ness.
        Outputs to object vector cell modules if they're separated,
        else to all abstract location modules
        Parameters
        -----------
            shiny: boolean location shiny-ness
        Returns
        -----------
            sigma_g: standard deviation of inferred abstract (grid cell) locations
        """
        return self.MLP_sigma_g_shiny(shiny)

    def f_sigma_p(self, p):
        """
        Use multi layer perceptron to generate standard deviation of grounded location retrieval.
        Parameters
        -----------
            p: grounded (place cell) locations
        Returns
        -----------
            sigma_p: standard deviation of grounded (place cell) locations
        """
        return self.MLP_sigma_p(p)

    def f_x(self, p):
        """
        Use multi layer perceptron to generate observation from grounded location. Do this by first calculating
        categorical probability distribution over observations for a given
        ground location, then sampling from that distribution.
        Parameters
        -----------
            p: grounded (place cell) locations
        Returns
        -----------
            x: generated observation
            logits: logits for generated observation
        """
        # Calculate categorical probability distribution over observations for a given ground location
        # p is the flattened (by concatenating rows - like reading sentences) outer product of g and x (p = g^T * x).
        # Therefore to get the sensory experience x for a grounded location p, sum over all abstract locations
        # g for each component of x. That's what the paper means when it says "sum over entorhinal preferences".
        # It can be done with the transpose of W_tile
        x = self.w_x * torch.matmul(p, torch.t(self.hyper["W_tile"][0])) + self.b_x
        # Then we need to decompress the temporally filtered sensory experience into a single current experience prediction
        logits = self.f_c_star(x)
        # We'll keep both the logits (domain -inf, inf) and probabilities (domain 0, 1) because both are needed later on
        probability = utils.softmax(logits)
        return probability, logits

    def f_c_star(self, compressed):
        """
        Use multi layer perceptron to decompress sensory experience at highest frequency
        Parameters
        -----------
            compressed: compressed sensory experience
        Returns
        -----------
            decompressed: decompressed sensory experience
        """
        return self.MLP_c_star(compressed)

    def f_c(self, decompressed):
        """
        Use multi layer perceptron to compress sensory experience at highest frequency
        Parameters
        -----------
            decompressed: decompressed sensory experience
        Returns
        -----------
            compressed: compressed sensory experience
        """
        # Compress sensory observation from one-hot provided by world to two-hot for ease of computation
        return torch.stack([self.hyper["two_hot_table"][i] for i in torch.argmax(decompressed, dim=1)], dim=0)

    def f_n(self, x):
        """
        Normalise sensory observation for each frequency module.
        Parameters
        -----------
            x: sensory observation, one-hot vector
        Returns
        -----------
            normalised: normalised sensory observation
        """
        normalised = [utils.normalise(utils.relu(x[f] - torch.mean(x[f]))) for f in range(self.hyper["n_f"])]
        return normalised

    def f_g(self, g):
        """
        Downsample abstract location for each frequency module.
        Parameters
        -----------
            g: abstract (grid cell) locations
        Returns
        -----------
            downsampled: downsampled abstract (grid cell) locations
        """
        downsampled = [torch.matmul(g[f], self.hyper["g_downsample"][f]) for f in range(self.hyper["n_f"])]
        return downsampled

    def f_g_clamp(self, g):
        """
        Calculate activation for abstract location, thresholding between -1 and 1.
        Parameters
        -----------
            g: abstract (grid cell) locations
        Returns
        -----------
            activation: activation for abstract (grid cell) locations
        """
        activation = [torch.clamp(g_f, min=-1, max=1) for g_f in g]
        return activation

    def f_p(self, p):
        """
        Calculate activation for inferred grounded location, using a leaky relu for sparsity. Either apply to full
        multi-frequency grounded location or single frequency module.
        Parameters
        -----------
            p: grounded (place cell) locations
        Returns
        -----------
            activation: activation for grounded (place cell) locations
        """
        activation = (
            [utils.leaky_relu(torch.clamp(p_f, min=-1, max=1)) for p_f in p]
            if type(p) is list
            else utils.leaky_relu(torch.clamp(p, min=-1, max=1))
        )
        return activation

    def attractor(self, p_query, M, retrieve_it_mask=None):
        """
        Retrieve grounded location from attractor network memory with weights M by pattern-completing query. For example,
        initial attractor input can come from abstract location
        (g_) or sensory experience (x_).
        Parameters
        -----------
            p_query: initial attractor input
            M: memory connectivity matrix
            retrieve_it_mask: mask for hierarchical retrieval
        Returns
        -----------
            p: retrieved grounded (place cell) locations
        """
        # Start by flattening query grounded locations across frequency modules
        h_t = torch.cat(p_query, dim=1)
        # Apply activation function to initial memory index
        h_t = self.f_p(h_t)
        # Hierarchical retrieval (not in paper) is implemented by early stopping retrieval for low frequencies,
        # using a mask. If not specified: initialise mask as all 1s
        retrieve_it_mask = (
            [torch.ones(sum(self.hyper["n_p"])) for _ in range(self.hyper["n_p"])]
            if retrieve_it_mask is None
            else retrieve_it_mask
        )
        # Iterate attractor dynamics to do pattern completion
        for tau in range(self.hyper["i_attractor"]):
            # Apply one iteration of attractor dynamics, but only where there is a 1 in the mask. NB retrieve_it_mask
            # entries have only one row, but are broadcasted to batch_size
            h_t = (1 - retrieve_it_mask[tau]) * h_t + retrieve_it_mask[tau] * (
                self.f_p(self.hyper["kappa"] * h_t + torch.squeeze(torch.matmul(torch.unsqueeze(h_t, 1), M)))
            )
        # Make helper list of cumulative neurons per frequency module for grounded locations
        n_p = np.cumsum(np.concatenate(([0], self.hyper["n_p"])))
        # Now re-cast the grounded location into different frequency modules, since memory retrieval turned it
        # into one long vector
        p = [h_t[:, n_p[f] : n_p[f + 1]] for f in range(self.hyper["n_f"])]
        return p

    def hebbian(self, M_prev, p_inferred, p_generated, do_hierarchical_connections=True):
        """
        Update attractor network memory by Hebbian learning of pattern. For example, initial attractor input can
        come from abstract location (g_) or sensory experience (x_).
        Parameters
        -----------
            M_prev: previous memory connectivity matrix
            p_inferred: inferred grounded (place cell) locations
            p_generated: generated grounded (place cell) locations
            do_hierarchical_connections: whether to use hierarchical connections
        Returns
        -----------
            M: updated memory connectivity matrix
        """
        # Create new ground memory for attractor network by setting weights to outer product of learned vectors
        # p_inferred corresponds to p in the paper, and p_generated corresponds to p^.
        # The order of p + p^ and p - p^ is reversed since these are row vectors, instead of column vectors in the paper.
        M_new = torch.squeeze(
            torch.matmul(torch.unsqueeze(p_inferred + p_generated, 2), torch.unsqueeze(p_inferred - p_generated, 1))
        )
        # Multiply by connection vector, e.g. only keeping weights from low to high frequencies for hierarchical retrieval
        if do_hierarchical_connections:
            M_new = M_new * self.hyper["p_update_mask"]
        # Store grounded location in attractor network memory with weights M by Hebbian learning of pattern
        M = torch.clamp(self.hyper["lambda"] * M_prev + self.hyper["eta"] * M_new, min=-1, max=1)
        return M


class MLP(torch.nn.Module):
    """
    Class for multi layer perceptron with multiple modules. Each module has two layers: input->hidden and hidden->output.
    """

    def __init__(self, in_dim, out_dim, activation=(torch.nn.functional.elu, None), hidden_dim=None, bias=(True, True)):
        """
        Initialise multi layer perceptron with multiple modules.
        """
        # First call super class init function to set up torch.nn.Module style model and inherit it's functionality
        super(MLP, self).__init__()
        # Check if this network consists of module: are input and output dimensions lists? If not, make them
        # (but remember it wasn't)
        if type(in_dim) is list:
            self.is_list = True
        else:
            in_dim = [in_dim]
            out_dim = [out_dim]
            self.is_list = False
        # Find number of modules
        self.N = len(in_dim)
        # Create weights (input->hidden, hidden->output) for each module
        self.w = torch.nn.ModuleList([])
        for n in range(self.N):
            # If number of hidden dimensions is not specified: mean of input and output
            if hidden_dim is None:
                hidden = int(np.mean([in_dim[n], out_dim[n]]))
            else:
                hidden = hidden_dim[n] if self.is_list else hidden_dim
                # Each module has two sets of weights: input->hidden and hidden->output
            self.w.append(
                torch.nn.ModuleList(
                    [torch.nn.Linear(in_dim[n], hidden, bias=bias[0]), torch.nn.Linear(hidden, out_dim[n], bias=bias[1])]
                )
            )
        # Copy activation function for hidden layer and output layer
        self.activation = activation
        # Initialise all weights
        with torch.no_grad():
            for from_layer in range(2):
                for n in range(self.N):
                    # Set weights to xavier initalisation
                    torch.nn.init.xavier_normal_(self.w[n][from_layer].weight)
                    # Set biases to 0
                    if bias[from_layer]:
                        self.w[n][from_layer].bias.fill_(0.0)

    def set_weights(self, from_layer, value):
        """
        Set weights for each module from a given layer to a given value.
        Parameters
        -----------
            from_layer: layer from which to set weights
            value: value to set weights to
        """
        # If single value is provided: copy it for each module
        if type(value) is not list:
            input_value = [value for n in range(self.N)]
        else:
            input_value = value
        # Run through all modules and set weights starting from requested layer to the specified value
        with torch.no_grad():
            # MLP is setup as follows: w[module][layer] is Linear object, w[module][layer].weight is Parameter object for
            # linear weights, w[module][layer].weight.data is tensor of weight values
            for n in range(self.N):
                # If a tensor is provided: copy the tensor to the weights
                if type(input_value[n]) is torch.Tensor:
                    self.w[n][from_layer].weight.copy_(input_value[n])
                    # If only a single value is provided: set that value everywhere
                else:
                    self.w[n][from_layer].weight.fill_(input_value[n])

    def forward(self, data):
        """
        Run input data through MLP network.
        Parameters
        -----------
            data: input data
        Returns
        -----------
            output: output data
        """
        # Make input data into list, if this network doesn't consist of modules
        if self.is_list:
            input_data = data
        else:
            input_data = [data]
        # Run input through network for each module
        output = []
        for n in range(self.N):
            # Pass through first weights from input to hidden layer
            module_output = self.w[n][0](input_data[n])
            # Apply hidden layer activation
            if self.activation[0] is not None:
                module_output = self.activation[0](module_output)
            # Pass through second weights from hidden to output layer
            module_output = self.w[n][1](module_output)
            # Apply output layer activation
            if self.activation[1] is not None:
                module_output = self.activation[1](module_output)
            # Transpose output again to go back to column vectors instead of row vectors
            output.append(module_output)
            # If this network doesn't consist of modules: select output from first module to return
        if not self.is_list:
            output = output[0]
        # And return output
        return output


class LSTM(torch.nn.Module):
    """
    Class for LSTM with multiple layers.
    """

    def __init__(self, in_dim, hidden_dim, out_dim, n_layers=1, n_a=4):
        """
        Initialise LSTM with multiple layers.
        """
        # First call super class init function to set up torch.nn.Module style model and inherit it's functionality
        super(LSTM, self).__init__()
        # LSTM layer
        self.lstm = torch.nn.LSTM(in_dim, hidden_dim, n_layers, batch_first=True)
        # Hidden to output
        self.lin = torch.nn.Linear(hidden_dim, out_dim)
        # Copy number of actions, will be needed for input data vector
        self.n_a = n_a

    def forward(self, data, prev_hidden=None):
        """
        Run input data through LSTM network.
        Parameters
        -----------
            data: input data
            prev_hidden: previous hidden state
        Returns
        -----------
            out: output data
            lstm_hidden: hidden state
        """
        # If previous hidden and cell state are not provided: initialise them randomly
        if prev_hidden is None:
            hidden_state = torch.randn(self.lstm.num_layers, data.shape[0], self.lstm.hidden_size)
            cell_state = torch.randn(self.lstm.num_layers, data.shape[0], self.lstm.hidden_size)
            prev_hidden = (hidden_state, cell_state)
        # Run input through lstm
        lstm_out, lstm_hidden = self.lstm(data, prev_hidden)
        # Apply linear network to lstm output to get output: prediction at each timestep
        lin_out = self.lin(lstm_out)
        # And since we want a one-hot prediciton: do softmax on top
        out = utils.softmax(lin_out)
        # Return output and hidden state
        return out, lstm_hidden

    def prepare_data(self, data_in):
        """
        Prepare input data for LSTM network.
        Parameters
        -----------
            data_in: input data
        Returns
        -----------
            data: prepared input data
        """
        # Transform list of actions of each step into batch of one-hot row vectors
        actions = [
            torch.zeros((len(step[2]), self.n_a)).scatter_(1, torch.tensor(step[2]).unsqueeze(1), 1.0) for step in data_in
        ]
        # Concatenate observation and action together along column direction in each step
        vectors = [torch.cat((step[1], action), dim=1) for step, action in zip(data_in, actions)]
        # Then stack all these together along the second dimension, which is sequence length
        data = torch.stack(vectors, dim=1)
        # Return data in [batch_size, seq_len, input_dim] dimension as expected by lstm
        return data


class Iteration:
    """
    Class for storing all data from one iteration of the model.
    """

    def __init__(
        self,
        g=None,
        x=None,
        a=None,
        L=None,
        M=None,
        g_gen=None,
        p_gen=None,
        x_gen=None,
        x_logits=None,
        x_inf=None,
        g_inf=None,
        p_inf=None,
    ):
        """
        Initialise iteration with all data from one iteration of the model.
        """
        # Copy all inputs
        self.g = g
        self.x = x
        self.a = a
        self.L = L
        self.M = M
        self.g_gen = g_gen
        self.p_gen = p_gen
        self.x_gen = x_gen
        self.x_logits = x_logits
        self.x_inf = x_inf
        self.g_inf = g_inf
        self.p_inf = p_inf

    def correct(self):
        """
        Calculate accuracy of model prediction for this iteration.
        Returns
        -----------
            accuracy: accuracy of model prediction for this iteration
        """
        # Detach observation and all predictions
        observation = self.x.detach().numpy()
        predictions = [tensor.detach().numpy() for tensor in self.x_gen]
        # Did the model predict the right observation in this iteration?
        accuracy = [np.argmax(prediction, axis=-1) == np.argmax(observation, axis=-1) for prediction in predictions]
        return accuracy

    def detach(self):
        """
        Detach all tensors contained in this iteration.
        Returns
        -----------
            self: iteration with all tensors detached
        """
        # Detach all tensors contained in this iteration
        self.L = [tensor.detach() for tensor in self.L]
        self.M = [tensor.detach() for tensor in self.M]
        self.g_gen = [tensor.detach() for tensor in self.g_gen]
        self.p_gen = [tensor.detach() for tensor in self.p_gen]
        self.x_gen = [tensor.detach() for tensor in self.x_gen]
        self.x_inf = [tensor.detach() for tensor in self.x_inf]
        self.g_inf = [tensor.detach() for tensor in self.g_inf]
        self.p_inf = [tensor.detach() for tensor in self.p_inf]
        # Return self after detaching everything
        return self
