# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/30_rcur.ipynb.

# %% auto 0
__all__ = ['RecurrentAutoEncoderMixin', 'Encoder', 'Decoder', 'AutoEncoder']

# %% ../nbs/30_rcur.ipynb 6
import logging

# %% ../nbs/30_rcur.ipynb 8
#| export

# %% ../nbs/30_rcur.ipynb 11
#| export

# %% ../nbs/30_rcur.ipynb 13
try: import numpy as np
except ImportError: ...

# %% ../nbs/30_rcur.ipynb 15
try: import torch, torch.nn as nn
except ImportError: ...

# %% ../nbs/30_rcur.ipynb 17
#| export


# %% ../nbs/30_rcur.ipynb 19
from atyp import IntQ, Tensor, TensorQ, DeviceQ, DTypeQ
from chck import isint, isnone, notnone
from fpos import val1st
from putl import dropkeys, deconstruct

# %% ../nbs/30_rcur.ipynb 21
from .atyp import HiddenState, CellState, HiddenStates, HiddenStatesQ, RecurrentStatesQ
from .util import as4d, batches, channels
from .logr import EncoderLogMessage, DecoderLogMessage
from .enum import Channels, NonLinearity, RecurrentLayer, InitMethod
from .init import directions, directed_layers, init_r0, init_rn
from .kwds import recurrent_kwds, recurrent_encoder_kwds, recurrent_decoder_kwds

# %% ../nbs/30_rcur.ipynb 23
class RecurrentAutoEncoderMixin:
    def init_params(
        self, 
        
        input_size: int, 
        hidden_size: int,
        num_layers: int = 1, 
        bias: bool = True,
        batch_first: bool = True, 
        dropout: float = 0.2, 
        bidirectional: bool = False,
        proj_size: int = 0, 
        nonlinearity: NonLinearity = NonLinearity.Tanh, 
        device: DeviceQ = None, 
        dtype: DTypeQ = None, 
        
        kind: RecurrentLayer = RecurrentLayer.LSTM, 
        
        decoder: bool = False,
        
        init: InitMethod = InitMethod.orthogonal,
        gain: float = np.sqrt(2),

        make_relu: bool = False,
        make_linear: bool = False,
        linear_size: IntQ = None,
        logger: logging.Logger = None,
    ):
        
        params = recurrent_kwds(
            input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, bias=bias, 
            batch_first=batch_first, dropout=dropout, bidirectional=bidirectional,
            proj_size=proj_size, nonlinearity=nonlinearity, device=device, dtype=dtype,
            kind=RecurrentLayer(kind).value
        )
        params = params if not decoder else recurrent_decoder_kwds(**params)
        
        self.params = dict(
            **params, kind=kind, decoder=decoder, init=init, gain=gain, 
            make_relu=make_relu, make_linear=make_linear, linear_size=linear_size, log=logger
        )
        for k, v in self.params.items(): setattr(self, k, v)

        self.rmod = RecurrentLayer(kind).get(**recurrent_encoder_kwds(**params))
        InitMethod(init).get(self.rmod.weight_ih_l0, gain=gain)
        InitMethod(init).get(self.rmod.weight_hh_l0, gain=gain)

        if make_relu: 
            self.relu = nn.ReLU()

        if make_linear:
            hsize = val1st(params, ('input_size', 'hidden_size'), default=0)
            osize = val1st(params, ('output_size', 'input_size', 'hidden_size'), default=0)
            self.full = nn.Linear(hsize, osize, bias=bias, device=device, dtype=dtype)
            InitMethod(init).get(self.full.weight, gain=gain)

    def prep_inputs(
        self, x: Tensor, h0: TensorQ = None, c0: TensorQ = None,
        hcell: IntQ = None, hsize: IntQ = None
    ) -> (Tensor, RecurrentStatesQ):
        
        hcell = hcell if notnone(hcell) else getattr(self, 'hidden_size', None)
        hsize = hsize if notnone(hsize) else getattr(self, 'proj_size', None)
        
        if notnone(hsize) and not (hsize > 0): 
            hsize = hcell
        
        bsize = batches(x, self.batch_first)
        nlays = directed_layers(self.num_layers, self.bidirectional)        
        device = self.device if notnone(self.device) else x.device
        return init_r0(
            x, h0, c0, nlays, hcell, hsize, bsize, device, batch_first = True, bidirectional = None, 
            kind = RecurrentLayer(self.kind).value
        )

# %% ../nbs/30_rcur.ipynb 25
class Encoder(nn.Module, RecurrentAutoEncoderMixin):
    def __init__(
        self,         
        input_size: int, hidden_size: int, num_layers: int = 1, bias: bool = True,
        batch_first: bool = True, dropout: float = 0.2, bidirectional: bool = False,
        proj_size: int = 0, nonlinearity: NonLinearity = NonLinearity.Tanh, 
        device: DeviceQ = None, dtype: DTypeQ = None, 
        kind: RecurrentLayer = RecurrentLayer.LSTM, 
        init: InitMethod = InitMethod.orthogonal,
        gain: float = np.sqrt(2),

        make_relu: bool = False,
        make_linear: bool = False,
        linear_size: IntQ = None,
        logger: logging.Logger = None,
    ):
        super(Encoder, self).__init__()
        self.init_params(
            input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, bias=bias, 
            batch_first=batch_first, dropout=dropout, bidirectional=bidirectional,
            proj_size=proj_size, nonlinearity=nonlinearity, device=device, dtype=dtype,
            kind = kind, decoder = False, init = init, gain =gain, 
            make_relu=make_relu, make_linear=make_linear, linear_size=linear_size, logger=logger
        )

    def latent(self, x, retlast: bool = False):
        self.eval()
        out, states = self.forward(x, retlast=retlast)
        hidden = states if self.is_gru() else states[0]
        latent = hidden.squeeze(0).detach()
        return latent

    def forward(self, x, retlast: bool = False):
        # set initial hidden and cell states
        x, (h0, c0) = self.prep_inputs(x)
        # forward propagate rmod
        EncoderLogMessage('post-prep', x.shape, h0.shape, c0.shape)(self.log)
        out, (h0, c0) = self.rmod(x, (h0, c0))
        # out: tensor of shape (batch_size, sequence_len, hidden_size) e.g. (769, 3, 12)
        # (h0, c0) both have [num_layers, batch_size, hidden_size]
        EncoderLogMessage('post-rmod', out.shape, h0.shape, c0.shape)(self.log)
        return init_rn(out, retlast), (h0, c0)

# %% ../nbs/30_rcur.ipynb 27
class Decoder(nn.Module, RecurrentAutoEncoderMixin):
    def __init__(
        self, 
        hidden_size: int, output_size: int, num_layers: int = 1, bias: bool = True,
        batch_first: bool = True, dropout: float = 0.2, bidirectional: bool = False,
        proj_size: int = 0, nonlinearity: NonLinearity = NonLinearity.Tanh, 
        device: DeviceQ = None, dtype: DTypeQ = None, 
        kind: RecurrentLayer = RecurrentLayer.LSTM, 
        init: InitMethod = InitMethod.orthogonal,
        gain: float = np.sqrt(2),

        make_relu: bool = False,
        make_linear: bool = True,
        linear_size: IntQ = None,        
        use_encoded_space: bool = True, logger: logging.Logger = None,
    ):
        super(Decoder, self).__init__()
        params = dict(
            input_size = hidden_size, hidden_size = output_size, num_layers = num_layers, bias = bias, 
            batch_first = batch_first, dropout = dropout, bidirectional = bidirectional, 
            proj_size = proj_size, nonlinearity = nonlinearity, device = device, 
            dtype = dtype, 
            kind = kind, decoder = True, init = init, gain =gain,
            make_relu=make_relu, make_linear=True, linear_size=linear_size, logger = logger
        )
        if use_encoded_space: params.update(hidden_size = hidden_size, linear_size = output_size)
        self.init_params(**params)
        self.use_encoded_space = use_encoded_space        
        
    def iterdec(self, out, **kwargs):
        res, h0, c0, seqlen, retlast = [], *deconstruct(kwargs, 'h0', 'c0', 'seqlen', 'retlast')

        for i in range(seqlen or 1):
            DecoderLogMessage('iter-step', out.shape, h0.shape, c0.shape, step_idx=i)(self.log)
            out, (h0, c0) = self.rmod(out, (h0, c0))
            res.append(out)

        res = torch.cat(res, dim=0)
        if self.use_encoded_space: res = self.full(res)
        DecoderLogMessage('post-iter', res.shape, h0.shape, c0.shape)(self.log)
        return init_rn(res, retlast), (h0, c0)
    
    def prep_states(self, x, h0 = None, c0 = None):
        if not self.use_encoded_space:
            if notnone(h0): h0 = self.full(h0) # map from encoder hidden size to decoder hidden size
            if notnone(c0): c0 = self.full(c0) # e.g. (nlayers, hidsize) --> (nlayers, outsize)
        
        out, (h0, c0) = self.prep_inputs(x, h0, c0, hcell=self.output_size, hsize=self.output_size)
        DecoderLogMessage('post-prep', out.shape, h0.shape, c0.shape)(self.log)
        return out, (h0, c0)
        
    def fulldec(self, out, **kwargs):
        h0, c0, retlast = deconstruct(kwargs, 'h0', 'c0', 'retlast')
        out, (h0, c0) = self.rmod(out, (h0, c0))
        DecoderLogMessage('post-rmod', out.shape, h0.shape, c0.shape)(self.log)
        return init_rn(out, retlast), (h0, c0)

    def forward(self, x, h0 = None, c0 = None, seqlen: int = None, retlast: bool = False):
        # set initial hidden and cell states
        if not isint(seqlen): seqlen = None
        out, (h0, c0) = self.prep_states(x, h0, c0)        
        if not self.use_encoded_space and seqlen is None:
            return self.fulldec(out, h0=h0, c0=c0, retlast=retlast)
        return self.iterdec(out, h0=h0, c0=c0, seqlen=seqlen, retlast=retlast)

# %% ../nbs/30_rcur.ipynb 29
class AutoEncoder(nn.Module, RecurrentAutoEncoderMixin):
    def __init__(
        self, 
        input_size: int, hidden_size: int, num_layers: int = 1, bias: bool = True,
        batch_first: bool = True, dropout: float = 0.2, bidirectional: bool = False,
        proj_size: int = 0, nonlinearity: NonLinearity = NonLinearity.Tanh, 
        device: DeviceQ = None, dtype: DTypeQ = None, 
        kind: RecurrentLayer = RecurrentLayer.LSTM, 
        init: InitMethod = InitMethod.orthogonal,
        gain: float = np.sqrt(2),

        make_relu: bool = False,
        make_linear: bool = False,
        linear_size: IntQ = None,        
        use_encoded_space: bool = True, logger: logging.Logger = None,
    ):
        super(AutoEncoder, self).__init__()      
        params = dict(
            input_size=input_size, hidden_size=hidden_size, output_size=input_size, 
            num_layers=num_layers, bias=bias, batch_first=batch_first, dropout=dropout, 
            bidirectional=bidirectional, proj_size=proj_size, nonlinearity=nonlinearity,
            device=device, dtype=dtype, 
            
            kind = kind, init = init, gain =gain,
            make_relu=make_relu, make_linear=make_linear, linear_size=linear_size, logger = logger,
            
        )
        for k, v in params.items(): setattr(self, k, v)
        
        drop = ('input_size', 'hidden_size', 'output_size', )
        self.encoder = Encoder(input_size, hidden_size, **dropkeys(params, drop))
        self.decoder = Decoder(hidden_size, input_size, **dropkeys(params, drop), use_encoded_space=use_encoded_space)
                
        self.use_encoded_space = use_encoded_space

    def latent(self, x, retlast: bool = False):
        return self.encoder.latent(x, retlast)

    def encode(self, x, retlast: bool = False):
        return self.encoder(x, retlast)
    
    def decode(self, x, h0 = None, c0 = None, seqlen: int = None, retlast: bool = False):
        return self.decoder(x, h0, c0, seqlen=seqlen, retlast=retlast)

    
    def forward(self, x, retlast: bool = False, width: IntQ = None, height: IntQ = None):
        # x.shape = [769, 3, 4096]
        seqlen = channels(x, self.batch_first) # 3
        enc_x, (h0, c0) = self.encoder(x, retlast=retlast) # encoded_x.shape = [769, 1, 12]
        seqlen = (seqlen or 1) if self.use_encoded_space else None
        
        dec_x, (h0, c0) = self.decoder(
            enc_x, h0, c0, seqlen=seqlen, retlast=retlast
        ) # decoded_x.shape = [769, nchannels, 4096]
        if notnone(width) and notnone(height):
            dec_x = as4d(dec_x, width, height, seqlen)
        return dec_x
