"""
Various models

Autoencoder etiquette:
- The LAST output should always be the encoded bottleneck layer (with the exception of DCA)
"""

import os
import sys
import logging
from typing import List, Tuple, Union, Callable
import functools

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import skorch
import skorch.utils

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import activations
import model_utils


torch.backends.cudnn.deterministic = True  # For reproducibility
torch.backends.cudnn.benchmark = False


class DeepCountAutoencoder(nn.Module):
    """Replicate of DCA"""

    def __init__(
        self,
        input_dim: int,
        inter_dim: int = 64,
        bottle_dim: int = 32,
        output_dim: int = None,
        mode: str = "zinb",
    ):
        super(DeepCountAutoencoder, self).__init__()
        torch.manual_seed(1234)  # Fixed initialization
        assert mode in ["zinb", "nb", "poisson"], f"Unrecognized mode: {mode}"
        self.mode = mode

        self.encoder = nn.Linear(input_dim, inter_dim)
        nn.init.xavier_uniform_(self.encoder.weight)  # Aka Glorot initialization
        self.bn1 = nn.BatchNorm1d(num_features=inter_dim)

        self.bottleneck = nn.Linear(inter_dim, bottle_dim)
        nn.init.xavier_uniform_(self.bottleneck.weight)
        self.bn2 = nn.BatchNorm1d(num_features=bottle_dim)

        self.decoder = nn.Linear(bottle_dim, inter_dim)
        nn.init.xavier_uniform_(self.decoder.weight)
        self.bn3 = nn.BatchNorm1d(num_features=inter_dim)

        output_dim = output_dim if output_dim is not None else input_dim
        self.mean = nn.Linear(
            inter_dim, output_dim
        )  # These parameters are all per-gene
        nn.init.xavier_uniform_(self.mean.weight)
        if "nb" in self.mode:
            self.disp = nn.Linear(inter_dim, output_dim)
            nn.init.xavier_uniform_(self.disp.weight)
        if self.mode == "zinb":
            self.dropout = nn.Linear(inter_dim, output_dim)
            nn.init.xavier_uniform_(self.dropout.weight)

        self.forward = self.forward_with_decode  # Default to returning denoised

    def encode(self, x):
        """Given input, return bottlenecked latent representation"""
        x = F.relu(self.bn1(self.encoder(x)))
        x = F.relu(self.bn2(self.bottleneck(x)))
        return x

    def decode(self, x, size_factors):
        """Given latent representation, output 3 output layers (mean, dispersion, dropout)"""
        x = F.relu(self.bn3(self.decoder(x)))
        mu = torch.clamp(torch.exp(self.mean(x)), min=1e-5, max=1e6)  # Mean
        sf_scaled = size_factors.view(-1, 1).repeat(1, mu.shape[1])
        mu_scaled = mu * sf_scaled  # Elementwise multiplication

        if "nb" in self.mode:
            theta = torch.clamp(
                F.softplus(self.disp(x)), min=1e-4, max=1e3
            )  # Dispersion
        else:  # Is poisson
            return mu_scaled

        if self.mode == "zinb":
            pi = torch.sigmoid(self.dropout(x))  # Dropout
            return mu_scaled, theta, pi
        else:
            return mu_scaled, theta

    def forward_with_decode(self, x, size_factors):
        """
        Return the parameters mu, theta, and pi

        The denoised matrix is generated by replacing the original count values with
        the mean of the negative binomial component as predicted in the output layer.
        This matrix represents the denoised and library size normalized expression matrix,
        the final output of the method.
        """
        return self.decode(self.encode(x), size_factors)

    def forward_no_decode(self, x, _size_factors):
        """
        Return the hidden representation instead of decoding it as well
        Useful for probing the latent dimension using pytorch_eval
        Use by setting model.forward = model.forward_no_decode
        """
        return self.encode(x)


class Encoder(nn.Module):
    def __init__(self, num_inputs: int, num_units=32, activation=nn.PReLU):
        super().__init__()
        self.num_inputs = num_inputs
        self.num_units = num_units

        self.encode1 = nn.Linear(self.num_inputs, 64)
        nn.init.xavier_uniform_(self.encode1.weight)
        self.bn1 = nn.BatchNorm1d(64)
        self.act1 = activation()

        self.encode2 = nn.Linear(64, self.num_units)
        nn.init.xavier_uniform_(self.encode2.weight)
        self.bn2 = nn.BatchNorm1d(num_units)
        self.act2 = activation()

    def forward(self, x):
        x = self.act1(self.bn1(self.encode1(x)))
        x = self.act2(self.bn2(self.encode2(x)))
        return x


class Decoder(nn.Module):
    def __init__(
        self,
        num_outputs: int,
        num_units: int = 32,
        intermediate_dim: int = 64,
        activation=nn.PReLU,
        final_activation=None,
    ):
        super().__init__()
        self.num_outputs = num_outputs
        self.num_units = num_units

        self.decode1 = nn.Linear(self.num_units, intermediate_dim)
        nn.init.xavier_uniform_(self.decode1.weight)
        self.bn1 = nn.BatchNorm1d(intermediate_dim)
        self.act1 = activation()

        self.decode21 = nn.Linear(intermediate_dim, self.num_outputs)
        nn.init.xavier_uniform_(self.decode21.weight)
        self.decode22 = nn.Linear(intermediate_dim, self.num_outputs)
        nn.init.xavier_uniform_(self.decode22.weight)
        self.decode23 = nn.Linear(intermediate_dim, self.num_outputs)
        nn.init.xavier_uniform_(self.decode23.weight)

        self.final_activations = nn.ModuleDict()
        if final_activation is not None:
            if isinstance(final_activation, list) or isinstance(
                final_activation, tuple
            ):
                assert len(final_activation) <= 3
                for i, act in enumerate(final_activation):
                    if act is None:
                        continue
                    self.final_activations[f"act{i+1}"] = act
            elif isinstance(final_activation, nn.Module):
                self.final_activations["act1"] = final_activation
            else:
                raise ValueError(
                    f"Unrecognized type for final_activation: {type(final_activation)}"
                )

    def forward(self, x, size_factors=None):
        """include size factor here because we may want to scale the output by that"""
        x = self.act1(self.bn1(self.decode1(x)))

        retval1 = self.decode21(x)  # This is invariably the counts
        if "act1" in self.final_activations.keys():
            retval1 = self.final_activations["act1"](retval1)
        if size_factors is not None:
            sf_scaled = size_factors.view(-1, 1).repeat(1, retval1.shape[1])
            retval1 = retval1 * sf_scaled  # Elementwise multiplication

        retval2 = self.decode22(x)
        if "act2" in self.final_activations.keys():
            retval2 = self.final_activations["act2"](retval2)

        retval3 = self.decode23(x)
        if "act3" in self.final_activations.keys():
            retval3 = self.final_activations["act3"](retval3)

        return retval1, retval2, retval3


class ShallowDecoder(nn.Module):
    """
    Shallow single-layer decoder to see if this might help with latent representation
    """

    def __init__(self, num_outputs: int, num_units: int = 32, final_activation=None):
        super().__init__()
        self.num_outputs = num_outputs
        self.num_units = num_units

        self.decode1 = nn.Linear(num_units, self.num_outputs)
        nn.init.xavier_uniform_(self.decode1.weight)
        self.decode2 = nn.Linear(num_units, self.num_outputs)
        nn.init.xavier_uniform_(self.decode2.weight)
        self.decode3 = nn.Linear(num_units, self.num_outputs)
        nn.init.xavier_uniform_(self.decode3.weight)

        self.final_activations = nn.ModuleDict()
        if final_activation is not None:
            if isinstance(final_activation, list) or isinstance(
                final_activation, tuple
            ):
                assert len(final_activation) <= 3
                for i, act in enumerate(final_activation):
                    if act is None:
                        continue
                    self.final_activations[f"act{i+1}"] = act
            elif isinstance(final_activation, nn.Module):
                self.final_activations["act1"] = final_activation
            else:
                raise ValueError(
                    f"Unrecognized type for final_activation: {type(final_activation)}"
                )

    def forward(self, x, size_factors=None):
        retval = self.decode1(x)
        if "act1" in self.final_activations.keys():
            retval1 = self.final_activations["act1"](retval1)
        if size_factors is not None:
            sf_scaled = size_factors.view(-1, 1).repeat(1, retval1.shape[1])
            retval1 = retval1 * sf_scaled  # Elementwise multiplication

        retval2 = self.decode2(x)
        if "act2" in self.final_activations.keys():
            retval2 = self.final_activations["act2"](retval2)

        retval3 = self.decode3(x)
        if "act3" in self.final_activations.keys():
            retval3 = self.final_activations["act3"](retval3)

        return retval1, retval2, retval3


class ChromEncoder(nn.Module):
    """
    Consumes multiple inputs (i.e. one feature vector for each chromosome)
    After processing everything to be the same dimensionality, concatenate
    to form a single latent dimension
    """

    def __init__(
        self, num_inputs: List[int], latent_dim: int = 32, activation=nn.PReLU
    ):
        super(ChromEncoder, self).__init__()
        self.num_inputs = num_inputs
        self.act = activation

        self.initial_modules = nn.ModuleList()
        for n in self.num_inputs:
            assert isinstance(n, int)
            layer1 = nn.Linear(n, 32)
            nn.init.xavier_uniform_(layer1.weight)
            bn1 = nn.BatchNorm1d(32)
            act1 = self.act()
            layer2 = nn.Linear(32, 16)
            nn.init.xavier_uniform_(layer2.weight)
            bn2 = nn.BatchNorm1d(16)
            act2 = self.act()
            self.initial_modules.append(
                nn.ModuleList([layer1, bn1, act1, layer2, bn2, act2])
            )

        self.encode2 = nn.Linear(16 * len(self.num_inputs), latent_dim)
        nn.init.xavier_uniform_(self.encode2.weight)
        self.bn2 = nn.BatchNorm1d(latent_dim)
        self.act2 = self.act()

    def forward(self, x):
        assert len(x) == len(
            self.num_inputs
        ), f"Expected {len(self.num_inputs)} inputs but got {len(x)}"
        enc_chroms = []
        for init_mod, chrom_input in zip(self.initial_modules, x):
            for f in init_mod:
                chrom_input = f(chrom_input)
            enc_chroms.append(chrom_input)
        enc1 = torch.cat(
            enc_chroms, dim=1
        )  # Concatenate along the feature dim not batch dim
        enc2 = self.act2(self.bn2(self.encode2(enc1)))
        return enc2


class ChromDecoder(nn.Module):
    """
    Network that is per-chromosome aware, but does not does not output
    per-chromsome values, instead concatenating them into a single vector
    """

    def __init__(
        self,
        num_outputs: List[int],  # Per-chromosome list of output sizes
        latent_dim: int = 32,
        activation=nn.PReLU,
        final_activations=[activations.Exp(), activations.ClippedSoftplus()],
    ):
        super().__init__()
        self.num_outputs = num_outputs
        self.latent_dim = latent_dim

        self.decode1 = nn.Linear(self.latent_dim, len(self.num_outputs) * 16)
        nn.init.xavier_uniform_(self.decode1.weight)
        self.bn1 = nn.BatchNorm1d(len(self.num_outputs) * 16)
        self.act1 = activation()

        self.final_activations = nn.ModuleDict()
        if final_activations is not None:
            if isinstance(final_activations, list) or isinstance(
                final_activations, tuple
            ):
                assert len(final_activations) <= 3
                for i, act in enumerate(final_activations):
                    if act is None:
                        continue
                    self.final_activations[f"act{i+1}"] = act
            elif isinstance(final_activations, nn.Module):
                self.final_activations["act1"] = final_activations
            else:
                raise ValueError(
                    f"Unrecognized type for final_activation: {type(final_activation)}"
                )
        logging.info(
            f"ChromDecoder with {len(self.final_activations)} output activations"
        )

        self.final_decoders = nn.ModuleList()  # List[List[Module]]
        for n in self.num_outputs:
            layer0 = nn.Linear(16, 32)
            nn.init.xavier_uniform_(layer0.weight)
            bn0 = nn.BatchNorm1d(32)
            act0 = activation()
            # l = [layer0, bn0, act0]
            # for _i in range(len(self.final_activations)):
            #     fc_layer = nn.Linear(32, n)
            #     nn.init.xavier_uniform_(fc_layer.weight)
            #     l.append(fc_layer)
            # self.final_decoders.append(nn.ModuleList(l))
            layer1 = nn.Linear(32, n)
            nn.init.xavier_uniform_(layer1.weight)
            layer2 = nn.Linear(32, n)
            nn.init.xavier_uniform_(layer2.weight)
            layer3 = nn.Linear(32, n)
            nn.init.xavier_uniform_(layer3.weight)
            self.final_decoders.append(
                nn.ModuleList([layer0, bn0, act0, layer1, layer2, layer3])
            )

    def forward(self, x):
        x = self.act1(self.bn1(self.decode1(x)))
        # This is the reverse operation of cat
        x_chunked = torch.chunk(x, chunks=len(self.num_outputs), dim=1)

        retval1, retval2, retval3 = [], [], []
        for chunk, processors in zip(x_chunked, self.final_decoders):
            # Each processor is a list of 3 different decoders
            # decode1, bn1, act1, *output_decoders = processors
            decode1, bn1, act1, decode21, decode22, decode23 = processors
            chunk = act1(bn1(decode1(chunk)))
            temp1 = decode21(chunk)
            temp2 = decode22(chunk)
            temp3 = decode23(chunk)

            if "act1" in self.final_activations.keys():
                # temp1 = output_decoders[0](chunk)
                temp1 = self.final_activations["act1"](temp1)
                # retval1.append(temp1)
            if "act2" in self.final_activations.keys():
                # temp2 = output_decoders[1](chunk)
                temp2 = self.final_activations["act2"](temp2)
                # retval2.append(temp2)
            if "act3" in self.final_activations.keys():
                # temp3 = output_decoders[2](chunk)
                temp3 = self.final_activations["act3"](temp3)
                # retval3.append(temp3)
            retval1.append(temp1)
            retval2.append(temp2)
            retval3.append(temp3)
        retval1 = torch.cat(retval1, dim=1)
        retval2 = torch.cat(retval2, dim=1)
        retval3 = torch.cat(retval3, dim=1)
        return retval1, retval2, retval3

        # retval = []
        # if retval1:
        #     retval.append(torch.cat(retval1, dim=1))
        # if retval2:
        #     retval.append(torch.cat(retval2, dim=1))
        # if retval3:
        #     retval.append(torch.cat(retval3, dim=1))
        # return tuple(retval)


class AutoEncoder(nn.Module):
    """Vanilla autoencoder"""

    def __init__(
        self,
        num_inputs: int,
        num_units: int = 32,
        num_outputs: int = None,
        activation=nn.PReLU,
        final_activation=None,
        seed=8947,
        output_encoded: bool = True,
    ):
        super().__init__()
        torch.manual_seed(seed)
        self.output_encoded = output_encoded
        self.num_inputs = num_inputs
        self.num_outputs = self.num_inputs if num_outputs is None else num_outputs
        self.num_units = num_units

        self.encoder = Encoder(
            self.num_inputs, num_units=self.num_units, activation=activation
        )
        self.decoder = Decoder(
            self.num_outputs,
            num_units=self.num_units,
            activation=activation,
            final_activation=final_activation,
        )

    def forward(self, X):
        encoded = self.encoder(X)
        decoded = self.decoder(encoded)[
            0
        ]  # Drop the second output that corresponds to addtl parameters
        if self.output_encoded:
            return decoded, encoded
        return decoded


class ChromAutoEncoder(nn.Module):
    """
    Autoencoder that maps per chromsome input to per chromsome output
    """

    def __init__(
        self,
        num_inputs: List[int],
        latent_dim: int = 32,
        activation=nn.PReLU,
        final_activation=activations.Exp(),
        seed: int = 4444,
        output_encoded: bool = True,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.output_encoded = output_encoded
        self.num_inputs = num_inputs
        self.latent_dim = latent_dim
        self.act = activation
        self.final_act = final_activation

        self.encoder = ChromEncoder(
            num_inputs=num_inputs, latent_dim=self.latent_dim, activation=self.act
        )
        self.decoder = ChromDecoder(
            num_outputs=num_inputs,
            latent_dim=self.latent_dim,
            activation=self.act,
            final_activations=self.final_act,
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)[0]
        if self.output_encoded:
            return decoded, encoded
        return decoded


class ZINBChromAutoEncoder(nn.Module):
    """"""

    def __init__(
        self,
        num_inputs: List[int],
        latent_dim: int = 32,
        activation=nn.PReLU,
        final_activations=[activations.Exp(), nn.Softplus(), nn.Sigmoid()],
        seed=4444,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.num_inputs = num_inputs
        self.latent_dim = latent_dim
        self.act = activation
        self.final_activations = final_activations

        self.encoder = ChromEncoder(
            num_inputs=num_inputs, latent_dim=self.latent_dim, activation=self.act
        )
        self.decoder = ChromDecoder(
            num_outputs=num_inputs,
            latent_dim=self.latent_dim,
            activation=self.act,
            final_activations=self.final_activations,
        )

    def forward(self, x):
        encoded = self.encoder(x)
        mean, dispersion, dropout = self.decoder(encoded)
        return mean, dispersion, dropout, encoded


class NBChromAutoEncoder(nn.Module):
    """"""

    def __init__(
        self,
        num_inputs: List[int],
        num_outputs: List[int] = None,
        latent_dim: int = 32,
        activation=nn.PReLU,
        final_activations=[activations.Exp(), nn.Softplus()],
        seed=4444,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.num_inputs = num_inputs
        self.num_outputs = num_inputs if num_outputs is None else num_outputs
        self.latent_dim = latent_dim
        self.act = activation
        self.final_activations = final_activations

        self.encoder = ChromEncoder(
            num_inputs=self.num_inputs, latent_dim=self.latent_dim, activation=self.act
        )
        self.decoder = ChromDecoder(
            num_outputs=self.num_outputs,
            latent_dim=self.latent_dim,
            activation=self.act,
            final_activations=self.final_activations,
        )

    def forward(self, x):
        encoded = self.encoder(x)
        mean, dispersion, *_ = self.decoder(encoded)
        return mean, dispersion, encoded


class ZINBGenomeChromAutoEncoder(nn.Module):
    """
    Autoencoder that consumes genome-wide input and outputs per-chrom (e.g. RNA > ATAC)
    """

    def __init__(
        self,
        num_inputs: int,
        num_outputs: List[int],
        latent_dim: int = 32,
        activation=nn.PReLU,
        final_activations=[activations.Exp(), nn.Softplus(), nn.Sigmoid()],
        seed=4444,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.num_inputs = num_inputs
        self.num_outputs = num_outputs
        self.latent_dim = latent_dim
        self.act = activation
        self.final_activations = final_activations

        self.encoder = Encoder(
            num_inputs, num_units=self.latent_dim, activation=self.act
        )
        self.decoder = ChromDecoder(
            num_outputs=self.num_outputs,
            latent_dim=self.latent_dim,
            activation=self.act,
            final_activations=self.final_activations,
        )

    def forward(self, x, size_factors=None):
        encoded = self.encoder(x)
        mean, dispersion, dropout = self.decoder(encoded)
        return mean, dispersion, dropout, encoded


class NBChromGenomeAutoEncoder(nn.Module):
    """
    Autoencoder that consumes per-chrom input and outputs genome-wide (e.g. ATAC > RNA)
    """

    def __init__(
        self,
        num_inputs: List[int],
        num_outputs: int,
        latent_dim: int = 32,
        activation=nn.PReLU,
        final_activations=[activations.Exp(), nn.Softplus()],
        seed=4444,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.num_inputs = num_inputs
        self.num_outputs = num_outputs
        self.latent_dim = latent_dim
        self.act = activation
        self.final_activations = final_activations

        self.encoder = ChromEncoder(
            self.num_inputs, latent_dim=self.latent_dim, activation=self.act
        )
        self.decoder = Decoder(
            self.num_outputs,
            num_units=self.latent_dim,
            activation=self.act,
            final_activation=self.final_activations,
        )

    def forward(self, x, size_factors=None):
        encoded = self.encoder(x)
        mean, dispersion, _ = self.decoder(encoded, size_factors=size_factors)
        return mean, dispersion, encoded


class NBAutoEncoder(nn.Module):
    """
    Negative binomial autoencoder - outputs a mean and dispersion instead of a single values
    Designed to output normalized counts
    """

    def __init__(
        self,
        num_inputs: int,
        num_units: int = 32,
        num_outputs: int = None,
        activation=nn.PReLU,
        final_activation=[activations.Exp(), nn.Softplus()],
        seed=8482,
    ):
        """softplus is a smooth approximation to relu and constrains to be always be positive"""
        super().__init__()
        torch.manual_seed(seed)
        self.num_inputs = num_inputs
        self.num_outputs = self.num_inputs if num_outputs is None else num_outputs
        self.num_units = num_units

        self.encoder = Encoder(
            self.num_inputs, num_units=self.num_units, activation=activation
        )
        self.decoder = Decoder(
            self.num_outputs,
            num_units=self.num_units,
            activation=activation,
            final_activation=final_activation,
        )

    def forward(self, X, size_factors=None):
        encoded = self.encoder(X)
        mean, dispersion, _ = self.decoder(encoded, size_factors=size_factors)
        return mean, dispersion, encoded  # <- return a tuple of two values

    def from_encoded(self, encoded):
        mean, dispersion, _ = self.decoder(encoded, size_factors=None)
        return mean, dispersion, encoded


class ZINBAutoEncoder(nn.Module):
    """
    ZINB autoencoder - output a mean, dispersion, and dropout instead of single values
    Designed to output normalized counts
    """

    def __init__(
        self,
        num_inputs: int,
        num_units: int = 32,
        num_outputs: int = None,
        activation=nn.PReLU,
        final_activation=[activations.Exp(), nn.Softplus(), nn.Sigmoid()],
        seed=2828,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.num_inputs = num_inputs
        self.num_outputs = self.num_inputs if num_outputs is None else num_outputs
        self.num_units = num_units
        self.encoder = Encoder(
            self.num_inputs, num_units=self.num_units, activation=activation
        )
        self.decoder = Decoder(
            self.num_outputs,
            num_units=self.num_units,
            activation=activation,
            final_activation=final_activation,
        )

    def forward(self, x, size_factors=None):
        encoded = self.encoder(x)
        mean, dispersion, dropout = self.decoder(encoded, size_factors=size_factors)
        return mean, dispersion, dropout, encoded

    def from_encoded(self, encoded):
        mean, dispersion, dropout = self.decoder(encoded, size_factors=None)
        return mean, dispersion, dropout, encoded


class PairedAutoEncoder(nn.Module):
    """
    Paired autoencoder
    Supports cross-domain prediction by naively combining/swapping
    encoder and decoder
    """

    def __init__(self, model1, model2):
        super().__init__()
        self.model1 = model1
        self.model2 = model2

    def forward(self, x):
        """x is expected to be a tuple of two inputs"""
        x1, x2 = x
        y1 = self.model1(x1)
        y2 = self.model2(x2)
        return (y1, y2)

    def translate_1_to_2(self, encoded1):
        """Given data from domain 1 output domain 2"""
        output2 = self.model2.from_encoded(encoded1)
        return output2

    def translate_2_to_1(self, encoded2):
        """Given data fromd omain 2 output domain 1"""
        output1 = self.model1.from_encoded(encoded2)
        return output1


class InvertibleAutoEncoder(nn.Module):
    """
    Similar to paired autoencoder, but with an additional invertible network linking
    the latent dimensions

    An example invertible network might be RealNVP
    """

    def __init__(self, model1, model2, link_model):
        super().__init__()
        self.model1 = model1
        self.model2 = model2
        self.link_model = link_model

    def forward(self, x):
        x1, x2 = x  # Unpack, expects tuple from PairedDataset
        # Run through models
        y1 = self.model1(x1)
        y2 = self.model2(x2)
        # Run encoded representations through link model
        enc1 = y1[-1]
        enc2 = y2[-1]

        enc2_pred = self.link_model.forward(enc1, mode="direct")
        enc1_pred = self.link_model.forward(enc2, mode="inverse")

        return (y1, y2, (enc1_pred, enc2_pred))

    def translate_1_to_2(self, encoded1):
        """Using invertible layer, translate domain 1 to domain 2"""
        encoded2 = self.link_model.forward(encoded1, mode="direct")[0]
        return self.model2.from_encoded(encoded2)

    def translate_2_to_1(self, encoded2):
        """Using invertbile layer, translate domain 2's latent representation to domain 1"""
        encoded1 = self.link_model.forward(encoded2, mode="inverse")[0]
        return self.model1.from_encoded(encoded1)


class CattedAutoEncoder(nn.Module):
    """
    Autoencoder that concatenates inputs after a first layer
    """

    def __init__(
        self,
        num_inputs1: int,
        num_inputs2: int,
        num_units: int = 32,
        final_activations1=[activations.Exp(), nn.Softplus()],
        final_activations2=[activations.Exp(), nn.Softplus(), nn.Sigmoid()],
        seed=9387,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.dim1 = num_inputs1
        self.dim2 = num_inputs2

        self.encoder1 = nn.Linear(num_inputs1, 64, bias=True)  # Input1 > 64
        self.bn1 = nn.BatchNorm1d(64)
        self.act1 = nn.PReLU()
        self.encoder2 = nn.Linear(num_inputs2, 64, bias=True)  # Input2 > 64
        self.bn2 = nn.BatchNorm1d(64)
        self.act2 = nn.PReLU()

        self.encoder_concat = nn.Linear(128, 32, bias=True)  # 128 > 32
        self.bn_encoder_concat = nn.BatchNorm1d(32)
        self.act_encoder = nn.PReLU()

        self.decoder_concat = nn.Linear(32, 128, bias=True)  # 32 > 128
        self.bn_decoder_concat = nn.BatchNorm1d(128)
        self.act_decoder = nn.PReLU()

        self.decoder1 = nn.Linear(64, num_inputs1, bias=True)
        self.bn_decoder1 = nn.BatchNorm1d(num_inputs1)
        self.decoder2 = nn.Linear(64, num_inputs2, bias=True)
        self.bn_decoder2 = nn.BatchNorm1d(num_inputs2)

        self.decoder1_outputs = nn.ModuleList()
        for i, act in enumerate(final_activations1):
            self.decoder1_outputs.append(
                nn.Sequential(nn.Linear(64, num_inputs1, bias=True), act)
            )
        self.decoder2_outputs = nn.ModuleList()
        for i, act in enumerate(final_activations2):
            self.decoder2_outputs.append(
                nn.Sequential(nn.Linear(64, num_inputs2, bias=True), act)
            )

    def forward(self, x, size_factors=None):
        if size_factors is not None:
            raise NotImplemented
        x1, x2 = x  # Unpack
        encoded1 = self.act1(self.bn1(self.encoder1(x1)))
        encoded2 = self.act2(self.bn2(self.encoder2(x2)))
        encoded_cat = torch.cat([encoded1, encoded2], axis=1)

        bottle = self.act_encoder(
            self.bn_encoder_concat(self.encoder_concat(encoded_cat))
        )

        decoded_cat = self.act_decoder(
            self.bn_decoder_concat(self.decoder_concat(bottle))
        )  # 128
        decoded1 = decoded_cat[:, :64]  # Split 128 into 2 64s
        decoded2 = decoded_cat[:, 64:]
        assert decoded1.shape == decoded2.shape

        # Do each output
        output1 = tuple([n(decoded1) for n in self.decoder1_outputs] + [bottle])
        output2 = tuple([n(decoded2) for n in self.decoder2_outputs] + [bottle])
        return output1, output2


class SplicedAutoEncoder(nn.Module):
    """
    Spliced Autoencoder - where we have 4 parts (2 encoders, 2 decoders) that are all combined
    This does not work when you have chromsome split features
    """

    def __init__(
        self,
        input_dim1: int,
        input_dim2: int,
        hidden_dim: int = 16,
        final_activations1: list = [activations.Exp(), nn.Softplus()],
        final_activations2=[activations.Exp(), nn.Softplus(), nn.Sigmoid()],
        flat_mode: bool = False,  # Controls if we have to re-split inputs
        seed=182822,
    ):
        super().__init__()
        torch.manual_seed(seed)

        self.flat_mode = flat_mode
        self.input_dim1 = input_dim1
        self.input_dim2 = input_dim2
        self.num_outputs1 = (
            len(final_activations1)
            if isinstance(final_activations1, (list, set, tuple))
            else 1
        )
        self.num_outputs2 = (
            len(final_activations2)
            if isinstance(final_activations2, (list, set, tuple))
            else 1
        )

        self.encoder1 = Encoder(num_inputs=input_dim1, num_units=hidden_dim)
        self.encoder2 = Encoder(num_inputs=input_dim2, num_units=hidden_dim)

        self.decoder1 = Decoder(
            num_outputs=input_dim1,
            num_units=hidden_dim,
            final_activation=final_activations1,
        )
        self.decoder2 = Decoder(
            num_outputs=input_dim2,
            num_units=hidden_dim,
            final_activation=final_activations2,
        )

    def split_catted_input(self, x):
        """
        Split catted input data to expected sizes
        """
        return torch.split(x, [self.input_dim1, self.input_dim2], dim=-1)

    def _combine_output_and_encoded(self, decoded, encoded, num_outputs: int):
        """
        Combines the output and encoded in a single output
        """
        if num_outputs > 1:
            retval = *decoded, encoded
        else:
            if isinstance(decoded, tuple):
                decoded = decoded[0]
            retval = decoded, encoded
        assert isinstance(retval, (list, tuple))
        assert isinstance(
            retval[0], (torch.TensorType, torch.Tensor)
        ), f"Expected tensor but got {type(retval[0])}"
        return retval

    def forward_single(
        self, x, size_factors=None, in_domain: int = 1, out_domain: int = 1
    ):
        """
        Return output of a single domain combination
        """
        assert in_domain in [1, 2] and out_domain in [1, 2]

        encoder = self.encoder1 if in_domain == 1 else self.encoder2
        decoder = self.decoder1 if out_domain == 1 else self.decoder2
        num_non_latent_out = self.num_outputs1 if out_domain == 1 else self.num_outputs2

        encoded = encoder(x)
        decoded = decoder(encoded)
        return self._combine_output_and_encoded(decoded, encoded, num_non_latent_out)

    def forward(self, x, size_factors=None, mode: Union[None, Tuple[int, int]] = None):
        if self.flat_mode:
            x = self.split_catted_input(x)
        assert isinstance(x, (tuple, list))
        assert len(x) == 2, "There should be two inputs to spliced autoencoder"
        encoded1 = self.encoder1(x[0])
        encoded2 = self.encoder2(x[1])

        decoded11 = self.decoder1(encoded1)
        retval11 = self._combine_output_and_encoded(
            decoded11, encoded1, self.num_outputs1
        )
        decoded12 = self.decoder2(encoded1)
        retval12 = self._combine_output_and_encoded(
            decoded12, encoded1, self.num_outputs2
        )
        decoded22 = self.decoder2(encoded2)
        retval22 = self._combine_output_and_encoded(
            decoded22, encoded2, self.num_outputs2
        )
        decoded21 = self.decoder1(encoded2)
        retval21 = self._combine_output_and_encoded(
            decoded21, encoded2, self.num_outputs1
        )

        if mode is None:
            return retval11, retval12, retval21, retval22
        retval_dict = {
            (1, 1): retval11,
            (1, 2): retval12,
            (2, 1): retval21,
            (2, 2): retval22,
        }
        if mode not in retval_dict:
            raise ValueError(f"Invalid mode code: {mode}")
        return retval_dict[mode]


class NaiveSplicedAutoEncoder(SplicedAutoEncoder):
    """
    Naive "spliced" autoencoder that does not use shared branches and instead simply
    trains four separate models
    """

    def __init__(
        self,
        input_dim1: int,
        input_dim2: int,
        hidden_dim: int = 16,
        final_activations1: Union[Callable, List[Callable]] = [
            activations.Exp(),
            activations.ClippedSoftplus(),
        ],
        final_activations2: Union[Callable, List[Callable]] = nn.Sigmoid(),
        flat_mode: bool = True,  # Controls if we have to re-split inputs
        seed: int = 182822,
    ):
        nn.Module.__init__(self)
        torch.manual_seed(seed)

        self.flat_mode = flat_mode
        self.input_dim1 = input_dim1
        self.input_dim2 = input_dim2
        self.num_outputs1 = (
            len(final_activations1)
            if isinstance(final_activations1, (list, set, tuple))
            else 1
        )
        self.num_outputs2 = (
            len(final_activations2)
            if isinstance(final_activations2, (list, set, tuple))
            else 1
        )

        self.model11 = nn.ModuleList(
            [
                Encoder(num_inputs=input_dim1, num_units=hidden_dim),
                Decoder(
                    num_outputs=input_dim1,
                    num_units=hidden_dim,
                    final_activation=final_activations1,
                ),
            ]
        )
        self.model12 = nn.ModuleList(
            [
                Encoder(num_inputs=input_dim1, num_units=hidden_dim),
                ChromDecoder(
                    num_outputs=input_dim2,
                    latent_dim=hidden_dim,
                    final_activations=final_activations2,
                ),
            ]
        )
        self.model22 = nn.ModuleList(
            [
                ChromEncoder(num_inputs=input_dim2, latent_dim=hidden_dim),
                ChromDecoder(
                    num_outputs=input_dim2,
                    latent_dim=hidden_dim,
                    final_activations=final_activations2,
                ),
            ]
        )
        self.model21 = nn.ModuleList(
            [
                ChromEncoder(num_inputs=input_dim2, latent_dim=hidden_dim),
                Decoder(
                    num_outputs=input_dim1,
                    num_units=hidden_dim,
                    final_activation=final_activations1,
                ),
            ]
        )
        # Each of these has exactly two items
        self.modlist_dict = nn.ModuleDict(
            {
                "11": self.model11,
                "12": self.model12,
                "22": self.model22,
                "21": self.model21,
            }
        )

    def split_catted_input(self, x):
        """Split the input into chunks that goes to each input to model"""
        a, b = torch.split(x, [self.input_dim1, sum(self.input_dim2)], dim=-1)
        return (a, torch.split(b, self.input_dim2, dim=-1))

    def forward_single(
        self, x, size_factors=None, in_domain: int = 1, out_domain: int = 1
    ):
        """
        Return output of a single domain combination
        """
        num_non_latent_out = self.num_outputs1 if out_domain == 1 else self.num_outputs2
        modlist = self.modlist_dict[str(in_domain) + str(out_domain)]

        encoded = modlist[0](x)
        decoded = modlist[1](encoded)

        if num_non_latent_out > 1:
            retval = *decoded, encoded
        else:
            if isinstance(decoded, tuple):
                decoded = decoded[0]
            retval = decoded, encoded

        assert isinstance(retval, tuple)
        assert isinstance(retval[0], (torch.TensorType, torch.Tensor))
        return retval

    def forward(self, x, size_factors=None, mode: Union[None, Tuple[int, int]] = None):
        if self.flat_mode:
            x = self.split_catted_input(x)
        assert isinstance(x, (tuple, list))
        assert len(x) == 2, "There should be two inputs to spliced autoencoder"
        retval11 = self.forward_single(
            x[0], size_factors=size_factors, in_domain=1, out_domain=1
        )
        retval12 = self.forward_single(
            x[0], size_factors=size_factors, in_domain=1, out_domain=2
        )
        retval21 = self.forward_single(
            x[1], size_factors=size_factors, in_domain=2, out_domain=1
        )
        retval22 = self.forward_single(
            x[1], size_factors=size_factors, in_domain=2, out_domain=2
        )

        if mode is None:
            return retval11, retval12, retval21, retval22
        retval_dict = {
            (1, 1): retval11,
            (1, 2): retval12,
            (2, 1): retval21,
            (2, 2): retval22,
        }
        if mode not in retval_dict:
            raise ValueError(f"Invalid mode code: {mode}")
        return retval_dict[mode]


class AssymSplicedAutoEncoder(SplicedAutoEncoder):
    """
    Assymmetric spliced autoencoder where branch 2 is a chrom AE
    """

    def __init__(
        self,
        input_dim1: int,
        input_dim2: List[int],
        hidden_dim: int = 16,
        final_activations1: list = [activations.Exp(), activations.ClippedSoftplus()],
        final_activations2=nn.Sigmoid(),
        flat_mode: bool = True,  # Controls if we have to re-split inputs
        seed: int = 182822,
    ):
        # https://stackoverflow.com/questions/9575409/calling-parent-class-init-with-multiple-inheritance-whats-the-right-way
        nn.Module.__init__(self)
        torch.manual_seed(seed)

        self.flat_mode = flat_mode
        self.input_dim1 = input_dim1
        self.input_dim2 = input_dim2
        self.num_outputs1 = (
            len(final_activations1)
            if isinstance(final_activations1, (list, set, tuple))
            else 1
        )
        self.num_outputs2 = (
            len(final_activations2)
            if isinstance(final_activations2, (list, set, tuple))
            else 1
        )

        self.encoder1 = Encoder(num_inputs=input_dim1, num_units=hidden_dim)
        self.encoder2 = ChromEncoder(num_inputs=input_dim2, latent_dim=hidden_dim)

        self.decoder1 = Decoder(
            num_outputs=input_dim1,
            num_units=hidden_dim,
            final_activation=final_activations1,
        )
        self.decoder2 = ChromDecoder(
            num_outputs=input_dim2,
            latent_dim=hidden_dim,
            final_activations=final_activations2,
        )

    def split_catted_input(self, x):
        """Split the input into chunks that goes to each input to model"""
        a, b = torch.split(x, [self.input_dim1, sum(self.input_dim2)], dim=-1)
        return (a, torch.split(b, self.input_dim2, dim=-1))


class PerChromSplicedAutoEncoder(SplicedAutoEncoder):
    """
    Autoencoder that treats each chromsome completely separately
    Essentially a series of 22 independent autoencoders
    """

    def __init__(
        self,
        input_dim1: List[int],
        input_dim2: List[int],
        per_chrom_hidden_dim: int = 8,
        final_activations1: Union[Callable, List[Callable]] = [
            activations.Exp(),
            activations.ClippedSoftplus(),
        ],
        final_activations2: Union[Callable, List[Callable]] = nn.Sigmoid(),
        flat_mode: bool = True,
        seed: int = 348483,
    ):
        nn.Module.__init__(self)
        torch.manual_seed(seed)

        self.flat_mode = flat_mode
        self.input_dim1 = input_dim1
        self.input_dim2 = input_dim2
        assert len(self.input_dim1) == len(
            self.input_dim2
        ), f"Got different lengths for input dimensions: {len(self.input_dim1)} {len(self.input_dim2)}"

        self.num_outputs1 = (
            len(final_activations1)
            if isinstance(final_activations1, (list, tuple))
            else 1
        )
        self.num_outputs2 = (
            len(final_activations2)
            if isinstance(final_activations2, (list, tuple))
            else 1
        )

        self.encoders1 = nn.ModuleList(
            [
                Encoder(num_inputs=n, num_units=per_chrom_hidden_dim)
                for n in self.input_dim1
            ]
        )
        self.encoders2 = nn.ModuleList(
            [
                Encoder(num_inputs=n, num_units=per_chrom_hidden_dim)
                for n in self.input_dim2
            ]
        )

        self.decoders1 = nn.ModuleList(
            [
                Decoder(
                    num_outputs=n,
                    num_units=per_chrom_hidden_dim,
                    final_activation=final_activations1,
                )
                for n in self.input_dim1
            ]
        )
        self.decoders2 = nn.ModuleList(
            [
                Decoder(
                    num_outputs=n,
                    num_units=per_chrom_hidden_dim,
                    final_activation=final_activations2,
                )
                for n in self.input_dim2
            ]
        )

    def split_catted_input(self, x) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]:
        assert sum(self.input_dim1) + sum(self.input_dim2) == x.shape[1]
        x1, x2 = torch.split(x, [sum(self.input_dim1), sum(self.input_dim2)], dim=1)
        x1 = torch.split(x1, self.input_dim1, dim=1)
        x2 = torch.split(x2, self.input_dim2, dim=1)
        return x1, x2

    def forward_single(
        self, x, size_factors=None, in_domain: int = 1, out_domain: int = 1
    ):
        encoders = self.encoders1 if in_domain == 1 else self.encoders2
        decoders = self.decoders1 if out_domain == 1 else self.decoders2

        encoded = [enc(i) for enc, i in zip(encoders, x)]
        decoded = [dec(latent) for dec, latent in zip(decoders, encoded)]

        # Concatenate the per-chrom outputs
        decoded = [torch.cat(part, dim=1) for part in list(zip(*decoded))]

        # Attach the encoded space
        retval = decoded + [torch.cat(encoded, dim=1)]
        return retval


### SKORCH MODELS WRAPPERS GO HERE ###
class AutoEncoderSkorchNet(skorch.NeuralNet):
    """Subclassed so that we can easily extract the encoded layer"""

    def __init__(self, prop_reg_lambda: float = 0.0, *args, **kwargs):
        assert isinstance(prop_reg_lambda, float)
        super(AutoEncoderSkorchNet, self).__init__(*args, **kwargs)
        self.prop_reg_lambda = prop_reg_lambda

    def get_encoded_layer(self, x):
        """Return the encoded representation"""
        encoded = []
        for output in self.forward_iter(x, training=False):
            assert isinstance(output, tuple)
            e = output[-1]
            encoded.append(skorch.utils.to_numpy(e))
        return np.concatenate(encoded, 0)

    def get_prop_reg(self, y_pred, y_true):
        """
        Compute regularization based on the overall proportion of each gene in the batch
        """
        per_gene_counts = torch.sum(y_true, axis=0)  # Sum across the batch
        per_gene_counts_norm = per_gene_counts / torch.sum(per_gene_counts)

        per_gene_pred_counts = torch.sum(y_pred, axis=0)
        per_gene_pred_counts_norm = per_gene_pred_counts / torch.sum(
            per_gene_pred_counts
        )

        # L2 distance between the two
        # d = F.pairwise_distance(per_gene_pred_counts_norm, per_gene_counts_norm, p=2, eps=1e-6, keepdim=False)
        d2 = torch.pow(per_gene_counts_norm - per_gene_pred_counts_norm, 2).mean()
        # d = torch.sqrt(d2)
        return d2

    def get_prop_reg_pbulk(self, y_pred, pbulk):
        """
        Compute regularization based on pre-computed cluster-aggregated pseudobulk
        """
        return F.mse_loss(y_pred, pbulk)

    def get_loss(self, y_pred, y_true, X=None, training=False):
        """
        Adds in a regularization term
        If y_true is a tuple, then we assume that is it is (cell_truth, cell_cluster_psuedobulk_mean)

        Reference:
        https://github.com/skorch-dev/skorch/blob/4097e90/skorch/net.py
        """
        # print(type(y_pred), type(y_true))
        if isinstance(y_true, tuple) or isinstance(y_true, list):
            assert len(y_true) == 2
            y, y_cluster_pbulk = y_true
            y = skorch.utils.to_tensor(y, device=self.device)
            y_cluster_pbulk = skorch.utils.to_tensor(
                y_cluster_pbulk, device=self.device
            )
        else:
            y = skorch.utils.to_tensor(y_true, device=self.device)
        loss = self.criterion_(y_pred, y)

        # Add regularization term
        if self.prop_reg_lambda != 0.0:
            if isinstance(y_true, tuple) or isinstance(y_true, list):
                loss += self.prop_reg_lambda * self.get_prop_reg_pbulk(
                    y_pred[0], y_cluster_pbulk
                )
            else:
                loss += self.prop_reg_lambda * self.get_prop_reg(y_pred[0], y_true)

        # print(type(y_pred), type(y))
        return loss


class PairedAutoEncoderSkorchNet(skorch.NeuralNet):
    def forward_iter(self, X, training=False, device="cpu"):
        """Subclassed to work with tuples"""
        dataset = self.get_dataset(X)
        iterator = self.get_iterator(dataset, training=training)
        for data in iterator:
            Xi = skorch.dataset.unpack_data(data)[0]
            yp = self.evaluation_step(Xi, training=training)
            if isinstance(yp, tuple):
                yield model_utils.recursive_to_device(yp)  # <- modification here
            else:
                yield yp.to(device)

    def predict_proba(self, x):
        """Subclassed so calling predict produces a tuple of outputs"""
        y_probas1, y_probas2 = [], []
        for yp in self.forward_iter(x, training=False):
            assert isinstance(yp, tuple)
            yp1 = yp[0][0]
            yp2 = yp[1][0]
            y_probas1.append(skorch.utils.to_numpy(yp1))
            y_probas2.append(skorch.utils.to_numpy(yp2))
        y_proba1 = np.concatenate(y_probas1, 0)
        y_proba2 = np.concatenate(y_probas2, 0)
        return y_proba1, y_proba2

    def get_encoded_layer(self, x):
        """Get the encoded representation as a TUPLE of two elements"""
        encoded1, encoded2 = [], []
        for out1, out2, *_other in self.forward_iter(x, training=False):
            encoded1.append(out1[-1])
            encoded2.append(out2[-1])
        return np.concatenate(encoded1, axis=0), np.concatenate(encoded2, axis=0)

    def translate_1_to_2(self, x):
        enc1, enc2 = self.get_encoded_layer(x)
        device = next(self.module_.parameters()).device
        enc1_torch = torch.from_numpy(enc1).to(device)
        return self.module_.translate_1_to_2(enc1_torch)[0].detach().cpu().numpy()

    def translate_2_to_1(self, x):
        enc1, enc2 = self.get_encoded_layer(x)
        device = next(self.module_.parameters()).device
        enc2_torch = torch.from_numpy(enc2).to(device)
        return self.module_.translate_2_to_1(enc2_torch)[0].detach().cpu().numpy()


class SplicedAutoEncoderSkorchNet(PairedAutoEncoderSkorchNet):
    """
    Skorch wrapper for the SplicedAutoEncoder above.
    Mostly here to take care of how we calculate loss
    """

    def predict_proba(self, x):
        """
        Subclassed so that calling predict produces a tuple of 4 outputs
        """
        y_probas1, y_probas2, y_probas3, y_probas4 = [], [], [], []
        for yp in self.forward_iter(x, training=False):
            assert isinstance(yp, tuple)
            yp1 = yp[0][0]
            yp2 = yp[1][0]
            yp3 = yp[2][0]
            yp4 = yp[3][0]
            y_probas1.append(skorch.utils.to_numpy(yp1))
            y_probas2.append(skorch.utils.to_numpy(yp2))
            y_probas3.append(skorch.utils.to_numpy(yp3))
            y_probas4.append(skorch.utils.to_numpy(yp4))
        y_proba1 = np.concatenate(y_probas1)
        y_proba2 = np.concatenate(y_probas2)
        y_proba3 = np.concatenate(y_probas3)
        y_proba4 = np.concatenate(y_probas4)
        # Order: 1to1, 1to2, 2to1, 2to2
        return y_proba1, y_proba2, y_proba3, y_proba4

    def get_encoded_layer(self, x):
        """Get the encoded representation as a TUPLE of two elements"""
        encoded1, encoded2 = [], []
        for out11, out12, out21, out22 in self.forward_iter(x, training=False):
            encoded1.append(out11[-1])
            encoded2.append(out22[-1])
        return np.concatenate(encoded1, axis=0), np.concatenate(encoded2, axis=0)

    def translate_1_to_1(self, x) -> np.ndarray:
        retval = [
            skorch.utils.to_numpy(yp[0][0])
            for yp in self.forward_iter(x, training=False)
        ]
        return np.concatenate(retval)

    def translate_1_to_2(self, x) -> np.ndarray:
        retval = [
            skorch.utils.to_numpy(yp[1][0])
            for yp in self.forward_iter(x, training=False)
        ]
        return np.concatenate(retval)

    def translate_2_to_1(self, x) -> np.ndarray:
        retval = [
            skorch.utils.to_numpy(yp[2][0])
            for yp in self.forward_iter(x, training=False)
        ]
        return np.concatenate(retval)

    def translate_2_to_2(self, x) -> np.ndarray:
        retval = [
            skorch.utils.to_numpy(yp[3][0])
            for yp in self.forward_iter(x, training=False)
        ]
        return np.concatenate(retval)

    def score(self, true, pred):
        """
        Required for sklearn gridsearch
        Since sklearn uses convention of (true, pred) in its score functions
        We use the same
        https://scikit-learn.org/stable/modules/classes.html#module-sklearn.metrics
        """
        return self.get_loss(pred, true)


class PairedInvertibleAutoEncoderSkorchNet(PairedAutoEncoderSkorchNet):
    def translate_1_to_2(self, x):
        enc1, enc2 = self.get_encoded_layer(x)
        device = next(self.module_.parameters()).device
        enc1_torch = torch.from_numpy(enc1).to(device)
        return self.module_.translate_1_to_2(enc1_torch)[0].detach().cpu().numpy()

    def translate_2_to_1(self, x):
        enc1, enc2 = self.get_encoded_layer(x)
        device = next(self.module_.parameters()).device
        enc2_torch = torch.from_numpy(enc2).to(device)
        return self.module_.translate_2_to_1(enc2_torch)[0].detach().cpu().numpy()


if __name__ == "__main__":
    ae = AutoEncoder(1000)
    nbae = NBAutoEncoder(1000)

    chrom_encoder = NBChromAutoEncoder(num_inputs=[10, 20])
    print(chrom_encoder)
