# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_registry.ipynb.

# %% auto 0
__all__ = ['hn_relu', 'hn_repu5', 'hn_softmax', 'register_model', 'create_model', 'named_partial', 'simple_fwd', 'fwd_vec',
           'fwd_conv', 'hn', 'hn_relu_mnist', 'hn_relu_cifar', 'hn_repu5_mnist', 'hn_repu5_cifar', 'hn_softmax_mnist',
           'hn_softmax_cifar', 'conv_ham', 'conv_ham_avgpool_mnist', 'conv_ham_maxpool_mnist', 'conv_ham_avgpool_cifar',
           'conv_ham_maxpool_cifar', 'energy_attn', 'energy_attn_mnist', 'energy_attn_cifar']

# %% ../nbs/04_registry.ipynb 3
import hamux as hmx
from typing import *
import functools as ft
from fastcore.utils import *
import jax
import jax.numpy as jnp
import jax.tree_util as jtu
import treex as tx
from einops import rearrange

# %% ../nbs/04_registry.ipynb 7
__MODELS = {}

def register_model(fgen:Callable): # Function that returns a HAM with desired config
    """Register a function that returns a model configuration factory function.
    The name of the function acts as the retrieval key and must be unique across models"""
    __MODELS[fgen.__name__] = fgen
    return fgen

def create_model(mname:str, # Retrieve this stored model name
                 *args, # Passed to retrieved factory function
                 **kwargs): # Passed to retrieved factory function
    """Retrieve the model name from all registered models, passing `args` and `kwargs` to the factory function"""
    assert mname in __MODELS, f"Model '{mname}' has not been registered"
    return __MODELS[mname](*args, **kwargs)

def named_partial(f, *args, new_name=None, order=None, **kwargs):
    """Like `functools.partial` but also copies over function name and docstring. 
    
    If new_name is not None, use that as the name
    """
    fnew = ft.partial(f,*args,**kwargs)
    fnew.__doc__ = f.__doc__
    name = new_name if new_name is not None else f.__name__
    fnew.__name__ = name
    if order is not None: fnew.order=order
    elif hasattr(f,'order'): fnew.order=f.order
    return fnew

# %% ../nbs/04_registry.ipynb 13
def simple_fwd(model:hmx.HAM, # HAM where layer[0] is the image input and layer[-1] are the labels
               x: jnp.ndarray, # Starting point for clamped layer[0]
               depth: int, # Number of iterations for which to run the model
               dt: float, # Step size through time
               rng: Optional[jnp.ndarray]=None): # If provided, initialize states to random instead of 0
    """A simple version of the forward function for showing in the paper.

    All time constants `tau` are set to be 1 in our architecture, but this is variable
    """
    # Initialize hidden states to our image
    xs = model.init_states(x.shape[0], rng=rng)
    xs[0] = jnp.array(x)

    # Masks allow us to clamp our visible data over time
    masks = jtu.tree_map(lambda x: jnp.ones_like(x, dtype=jnp.int8), xs)
    masks[0] = jnp.zeros_like(masks[0], dtype=jnp.int8)  # Don't evolve images

    for i in range(depth):
        updates = model.vupdates(xs)  # Calculate the updates
        xs = model.step(
            xs, updates, dt=dt, masks=masks
        )  # Add them to our current states

    # All labels have a softmax activation function as the last layer, spitting out probabilities
    return model.layers[-1].g(xs[-1])

def fwd_vec(model:hmx.HAM, # HAM where layer[0] is the image input and layer[-1] are the labels
               x: jnp.ndarray, # Starting point for clamped layer[0]
               depth: int, # Number of iterations for which to run the model
               dt: float, # Step size through time
               rng: Optional[jnp.ndarray]=None): # If provided, initialize states to random instead of 0
    """Where the image input is vectorized"""
    x = rearrange(x, "... c h w -> ... (c h w)")
    return simple_fwd(model, x, depth, dt, rng)

def fwd_conv(model:hmx.HAM, # HAM where layer[0] is the image input and layer[-1] are the labels
               x: jnp.ndarray, # Starting point for clamped layer[0]
               depth: int, # Number of iterations for which to run the model
               dt: float, # Step size through time
               rng: Optional[jnp.ndarray]=None): # If provided, initialize states to random instead of 0
    """Where the image input is kept as a 3 channel image"""
    x = rearrange(x, "... c h w -> ... h w c")
    return simple_fwd(model,x, depth,dt, rng)


# %% ../nbs/04_registry.ipynb 16
@register_model
def hn(hidden_lagrangian:tx.Module,
       img_shape: Tuple, # Shape of image input to model
       label_shape: Tuple, # Shape of label probabilities,typically (NLABELS,)
       nhid:int=1000, # Number of units in hidden layer
       do_norm:bool=False): # If provided, enforce that all weights are standardized
    """Create a Classical Hopfield Network that is intended to be applied on vectorized inputs"""
    layers = [
        hmx.TanhLayer(img_shape),
        hmx.SoftmaxLayer(label_shape),
    ]

    synapses = [
        hmx.DenseMatrixSynapseWithHiddenLayer(nhid, hidden_lagrangian=hidden_lagrangian, do_norm=do_norm),
    ]

    connections = [
        ((0, 1), 0),
    ]

    ham = hmx.HAM(layers, synapses, connections)

    forward = ft.partial(fwd_vec, depth=4, dt=0.4)

    return ham, forward

hn_relu = named_partial(hn, hmx.lagrangians.LRelu(), new_name="hn_relu")
register_model(hn_relu)

hn_repu5 = named_partial(hn, hmx.lagrangians.LRepu(n=5), new_name="hn_repu5")
register_model(hn_repu5)

hn_softmax = named_partial(hn, hmx.lagrangians.LSoftmax(), new_name="hn_softmax")
register_model(hn_softmax)

@register_model
def hn_relu_mnist(nhid:int=1000): # Number of units in the single hidden layer
    """Vectorized HN on flattened MNIST"""
    return hn_relu(img_shape=(784,), label_shape=(10,), nhid=nhid)

@register_model
def hn_relu_cifar(nhid:int=6000): # Number of units in the single hidden layer
    """Vectorized HN on flattened CIFAR10"""
    return hn_relu(img_shape=(3072,), label_shape=(10,), nhid=nhid)

@register_model
def hn_repu5_mnist(nhid=1000):
    """Vectorized DAM on flattened MNIST"""
    return hn_repu5(img_shape=(784,), label_shape=(10,), nhid=nhid)

@register_model
def hn_repu5_cifar(nhid=6000):
    """Vectorized DAM on flattened CIFAR"""
    return hn_repu5(img_shape=(3072,), label_shape=(10,), nhid=nhid)

@register_model
def hn_softmax_mnist(nhid=1000):
    return hn_softmax(img_shape=(784,), label_shape=(10,), nhid=nhid, do_norm=True)

@register_model
def hn_softmax_cifar(nhid=6000):
    return hn_softmax(img_shape=(3072,), label_shape=(10,), nhid=nhid, do_norm=True)


# %% ../nbs/04_registry.ipynb 21
@register_model
def conv_ham(s1, s2, s3, pool_type, nhid=1000):
    layers = [
        hmx.TanhLayer(s1, tau=1.0),
        hmx.TanhLayer(s2, tau=1.0),
        hmx.TanhLayer(s3, tau=1.0),
        hmx.SoftmaxLayer((10,), tau=1.0),
    ]
    synapses = [
        hmx.ConvSynapseWithPool(
            (4, 4),
            strides=(2, 2),
            padding=(2, 2),
            pool_window=(2, 2),
            pool_stride=(2, 2),
            pool_type=pool_type,
        ),
        hmx.ConvSynapseWithPool(
            (3, 3),
            strides=(1, 1),
            padding=(0, 0),
            pool_window=(2, 2),
            pool_stride=(2, 2),
            pool_type=pool_type,
        ),
        hmx.DenseMatrixSynapseWithHiddenLayer(nhid),
    ]
    connections = [
        ((0, 1), 0), 
        ((1, 2), 1), 
        ((2, 3), 2)
    ]

    ham = hmx.HAM(layers, synapses, connections)

    forward = ft.partial(fwd_conv, depth=7, dt=0.3)
    return ham, forward


@register_model
def conv_ham_avgpool_mnist(nhid=1000):
    return conv_ham((28, 28, 1), (7, 7, 64), (2, 2, 128), pool_type="avg", nhid=nhid)


@register_model
def conv_ham_maxpool_mnist(nhid=1000):
    return conv_ham((28, 28, 1), (7, 7, 64), (2, 2, 128), pool_type="max", nhid=nhid)


@register_model
def conv_ham_avgpool_cifar(nhid=1000):
    return conv_ham((32, 32, 3), (8, 8, 90), (3, 3, 180), pool_type="avg", nhid=nhid)


@register_model
def conv_ham_maxpool_cifar(nhid=1000):
    return conv_ham((32, 32, 3), (8, 8, 90), (3, 3, 180), pool_type="max", nhid=nhid)

# %% ../nbs/04_registry.ipynb 24
@register_model
def energy_attn(s1, s2, nheads_self, nheads_cross):
    layers = [
        hmx.TanhLayer(s1, tau=1.0),
        hmx.TanhLayer(s2, tau=1.0, use_bias=True),
        hmx.SoftmaxLayer((10,), tau=1.0),
    ]

    synapses = [
        hmx.ConvSynapse((4, 4), strides=(4, 4), padding=(0, 0)),
        hmx.AttentionSynapse(num_heads=nheads_cross, zspace_dim=64, stdinit=0.002),
        hmx.AttentionSynapse(num_heads=nheads_self, zspace_dim=64, stdinit=0.002),
    ]

    connections = [(
        (0, 1), 0), 
        ((2, 1), 1), 
        ((1, 1), 2)
    ]
    ham = hmx.HAM(layers, synapses, connections)
    forward = ft.partial(fwd_conv, depth=5, dt=0.4)

    return ham, forward


@register_model
def energy_attn_mnist():
    return energy_attn(
        (28, 28, 1),
        (7, 7, 128),
        nheads_self=4,
        nheads_cross=2,
    )


@register_model
def energy_attn_cifar():
    return energy_attn(
        (32, 32, 3),
        (8, 8, 224),
        nheads_self=4,
        nheads_cross=2,
    )
