

# Checked for newly added batch emdb in `batch.batch_emb`.
# The disentangler only looks (potentially) for CT and NCC in batch.y, so all such cases were checked.

import os, sys
import numpy as np
import torch
import torch.nn as nn
from . import gnn
import torch_geometric as pyg
from . import mlp
#from enum import Enum

class ModeArch:
    HEADXINT = '0.ModeArch'
    HEADXSPL = '1.ModeArch'
    HEADBOTH = '2.ModeArch'
    TWOSEP = '3.ModeArch'

class ArchInsertionPoint:
    NONE = '0.ArchInsertionPoint'
    BACKBONE = '1.ArchInsertionPoint'
    HEADINT = '2.ArchInsertionPoint'
    HEADSPL = '3.ArchInsertionPoint'


class GNNDisentangler(nn.Module):
    def __init__(self, kwargs_genmodel, str_mode_normalizex:str, maxsize_subgraph, dict_general_args, mode_headxint_headxspl_headboth_twosep,
                 gnn_list_dim_hidden, kwargs_sageconv, clipval_cov_noncentralnodes:float, dict_CTNNC_usage:dict, std_minval_finalclip:float, std_maxval_finalclip:float,
                 flag_use_layernorm:bool, flag_use_dropout:bool, flag_enable_batchtoken:bool):
        '''
        :param maxsize_subgraph: the max size of the subgraph returned by pyg's NeighLoader.
        :param str_mode_normalizex: in ['counts', 'logp1'], whether the GNN works on counts or log1p
        :param dict_general_args: a dict containint
            - num_celltypes
            - flag_use_int_u
            - flag_use_spl_u
        :param mode_headxint_headxspl_headboth_twosep:
            - headxint: means that only muxint is generated by the head, and `muxspl =  x_cnt - muxint`
            - headxspl: means that only muxspl is generated by the head, and `muxint =  x_cnt - muxspl`
            - headboth: means muxint and muxspl are create by two different heads, in which case p(x|x_int+x_spl) must be included in logp(.)
            - twosep: two separated brances are considere
                - gnn for x_spl
                - mlp fro x_int
        :param clipval_cov_noncentralnodes: the value by which the covariance for non-central nodes is clipped.
        :param dict_CTNNC_usage: a dictionary that specifies wheteher/how CT and NCC are used.
        :param std_minval_finalclip, std_maxval_finalclip: values for the final clip. Can also be used to freez the std to a fixed value.
        :param flag_enable_batchtoken: if set to True, the inferred `muxint` and `muxspl` are shifted by batch-specific tensors.
        '''

        super(GNNDisentangler, self).__init__()
        self.kwargs_genmodel = kwargs_genmodel
        self.num_celltypes = dict_general_args['num_celltypes']
        self.flag_use_int_u = dict_general_args['flag_use_int_u']
        self.flag_use_spl_u = dict_general_args['flag_use_spl_u']
        self.mode_headxint_headxspl_headboth_twosep = mode_headxint_headxspl_headboth_twosep
        self.gnn_list_dim_hidden = gnn_list_dim_hidden
        self.kwargs_sageconv = kwargs_sageconv
        self.str_mode_normalizex = str_mode_normalizex
        self.clipval_cov_noncentralnodes = clipval_cov_noncentralnodes
        self.dict_CTNNC_usage = dict_CTNNC_usage
        self.std_minval_finalclip = std_minval_finalclip
        self.std_maxval_finalclip = std_maxval_finalclip
        self.flag_use_layernorm = flag_use_layernorm
        self.flag_use_dropout = flag_use_dropout
        self.flag_enable_batchtoken = flag_enable_batchtoken

        self._check_dict_CTNCCusage()
        self._check_args()

        # define module_gnn, which is the backbone in the 3 cases, and appears in the spl branch in TWOSEP mode
        dim_gnnin = kwargs_genmodel['dict_varname_to_dim']['x']
        if self.dict_CTNNC_usage['CT'] == ArchInsertionPoint.BACKBONE:
            # assert self.mode_headxint_headxspl_headboth_twosep != ModeArch.TWOSEP  # twosep --> backbone not defined
            dim_gnnin += kwargs_genmodel['dict_varname_to_dim']['CT']

        if self.dict_CTNNC_usage['NCC'] == ArchInsertionPoint.BACKBONE:
            # assert self.mode_headxint_headxspl_headboth_twosep != ModeArch.TWOSEP  # twosep --> backbone not defined
            dim_gnnin += kwargs_genmodel['dict_varname_to_dim']['NCC']

        # add the batch token
        if self.flag_enable_batchtoken:
            dim_gnnin += kwargs_genmodel['dict_varname_to_dim']['BatchEmb']


        '''
        if self.mode_headxint_headxspl_headboth_twosep == ModeArch.TWOSEP:
            if self.dict_CTNNC_usage['CT'] == ArchInsertionPoint.HEADSPL:
                dim_gnnin += kwargs_genmodel['dict_varname_to_dim']['CT']
            if self.dict_CTNNC_usage['NCC'] == ArchInsertionPoint.HEADSPL:
                dim_gnnin += kwargs_genmodel['dict_varname_to_dim']['NCC']
        '''

        self.str_mode_headxint_headxspl_headboth = {
            ModeArch.HEADXINT:'headxint',
            ModeArch.HEADXSPL:'headxspl',
            ModeArch.HEADBOTH:'headboth',
            ModeArch.TWOSEP:'twosep'
        }[self.mode_headxint_headxspl_headboth_twosep]  # backward comtblity (to be used only by the outer modules)


        self.module_gnn = gnn.SAGE(
            dim_input=dim_gnnin,
            dim_output=100,  # TODO:TUNE 100
            list_dim_hidden=self.gnn_list_dim_hidden,
            kwargs_sageconv=self.kwargs_sageconv
        )

        # add the batch-specific shifts (to be applied on GNN's output)
        if self.flag_enable_batchtoken:
            if kwargs_genmodel['dict_varname_to_dim']['BatchEmb'] == 1:
                pass
                # raise NotImplementedError(
                #     "When having only one batch, in your `config_model.yml` file `flag_enable_batchtoken_disentangler` is required to be 'False'."
                # )

            if kwargs_genmodel['dict_varname_to_dim']['BatchEmb'] == 1:
                self.param_batchshifts_4_gnnoutput = None
            else:
                self.param_batchshifts_4_gnnoutput = torch.nn.Parameter(
                    torch.randn(
                        [kwargs_genmodel['dict_varname_to_dim']['BatchEmb']-1, self.module_gnn.dim_output],
                        requires_grad=True
                    ),
                    requires_grad=True
                )  # [num_batch-1 x gnnoutput]

        # assert that the GNN backbone has as many hops as the generative model
        cnt_sage_conv = 0
        for ch in self.module_gnn.list_modules:
            if isinstance(ch, pyg.nn.SAGEConv) or isinstance(ch, gnn.SageConvAndActivation):
                cnt_sage_conv += 1

        assert (
            cnt_sage_conv == self.kwargs_genmodel['kwargs_theta_aggr']['num_hops']
        )

        # determine adddim_int_mlp, adddim_spl_mlp (the extension by CT or NCC) ===
        # int
        adddim_int_mlp = 0
        if self.dict_CTNNC_usage['CT'] == ArchInsertionPoint.HEADINT:
            adddim_int_mlp += kwargs_genmodel['dict_varname_to_dim']['CT']
        if self.dict_CTNNC_usage['NCC'] == ArchInsertionPoint.HEADINT:
            adddim_int_mlp += kwargs_genmodel['dict_varname_to_dim']['NCC']
        # spl
        adddim_spl_mlp = 0
        if self.dict_CTNNC_usage['CT'] == ArchInsertionPoint.HEADSPL:
            adddim_spl_mlp += kwargs_genmodel['dict_varname_to_dim']['CT']
        if self.dict_CTNNC_usage['NCC'] == ArchInsertionPoint.HEADSPL:
            adddim_spl_mlp += kwargs_genmodel['dict_varname_to_dim']['NCC']


        # create the int and spl MLPs
        if self.mode_headxint_headxspl_headboth_twosep == ModeArch.HEADXINT:
            self.module_muxint = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_int_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x']//10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x']//10) if(self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_covxint = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_int_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_muxspl = None
            self.module_covxspl = None
        elif self.mode_headxint_headxspl_headboth_twosep == ModeArch.HEADXSPL:
            self.module_muxint = None
            self.module_covxint = None
            self.module_muxspl = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_spl_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_covxspl = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_spl_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10

        elif self.mode_headxint_headxspl_headboth_twosep == ModeArch.HEADBOTH:
            self.module_muxint = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_int_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_covxint = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_int_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_muxspl = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_spl_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_covxspl = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_spl_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
        elif self.mode_headxint_headxspl_headboth_twosep == ModeArch.TWOSEP:
            # in this case the MLP on intrinsic part directly takes in the gex vector, but the MLP on the spatial branch takes in gnn-s output
            self.module_muxint = nn.Sequential(
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] + adddim_int_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_covxint = nn.Sequential(
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] + adddim_int_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_muxspl = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_spl_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10
            self.module_covxspl = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p=0.1) if (self.flag_use_dropout) else mlp.Identity(),
                nn.Linear(
                    self.module_gnn.dim_output + adddim_spl_mlp,
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10
                ),
                nn.LayerNorm(kwargs_genmodel['dict_varname_to_dim']['x'] // 10) if (self.flag_use_layernorm) else mlp.Identity(),
                nn.ReLU(),
                nn.Linear(
                    kwargs_genmodel['dict_varname_to_dim']['x'] // 10,
                    kwargs_genmodel['dict_varname_to_dim']['x']
                )
            )  # TODO:TUNE //10

        else:
            raise Exception(
                "Uknown value {} for str_mode_headxint_headxspl_headboth.".format(self.mode_headxint_headxspl_headboth_twosep)
            )









    def _check_dict_CTNCCusage(self):
        # format of the dict =====
        assert (
            self.mode_headxint_headxspl_headboth_twosep in [ModeArch.HEADXINT, ModeArch.HEADXSPL, ModeArch.HEADBOTH, ModeArch.TWOSEP]
        )
        assert (
            set(self.dict_CTNNC_usage.keys()) == {'CT', 'NCC'}
        )
        for k in self.dict_CTNNC_usage:
            assert (
                self.dict_CTNNC_usage[k] in [ArchInsertionPoint.NONE, ArchInsertionPoint.BACKBONE, ArchInsertionPoint.HEADINT, ArchInsertionPoint.HEADSPL]
            )

        # logic of the dict ====
        # CT and NCC cannot be entered into x-xpl or x-xint brances
        if self.mode_headxint_headxspl_headboth_twosep == ModeArch.HEADXINT:
            assert (
                self.dict_CTNNC_usage['CT'] != ArchInsertionPoint.HEADSPL
            )
            assert (
                self.dict_CTNNC_usage['NCC'] != ArchInsertionPoint.HEADSPL
            )

        if self.mode_headxint_headxspl_headboth_twosep == ModeArch.HEADXSPL:
            assert (
                self.dict_CTNNC_usage['CT'] != ArchInsertionPoint.HEADINT
            )
            assert (
                self.dict_CTNNC_usage['NCC'] != ArchInsertionPoint.HEADINT
            )








    def _check_args(self):
        assert (
            self.flag_enable_batchtoken in [True, False]
        )
        assert (
            self.flag_use_dropout in [True, False]
        )
        assert (
            self.flag_use_layernorm in [True, False]
        )
        assert (
            self.mode_headxint_headxspl_headboth_twosep in [ModeArch.HEADXINT, ModeArch.HEADXSPL, ModeArch.HEADBOTH, ModeArch.TWOSEP]
        )
        assert (
            self.kwargs_genmodel['dict_varname_to_dim']['CT'] == self.kwargs_genmodel['dict_varname_to_dim']['NCC']
        )
        assert (
            self.str_mode_normalizex in ['counts', 'log1p']
        )

    def _feed_to_GNN(self, x_input, batch, device):
        '''
        Handles different cases of feeding to the GNN.
        :param x_input:
        :param batch:
        :return:
        '''
        ten_input_gnn = self._extend_input(
            ten_input=x_input,
            key_inspoint=ArchInsertionPoint.BACKBONE,
            batch=batch,
            device=device
        )

        # add the batch-token
        if self.flag_enable_batchtoken:
            rng_batchemb = [
                batch.INFLOWMETAINF['dim_u_int'] + batch.INFLOWMETAINF['dim_u_spl'] + batch.INFLOWMETAINF['dim_CT'] + batch.INFLOWMETAINF['dim_NCC'],
                batch.INFLOWMETAINF['dim_u_int'] + batch.INFLOWMETAINF['dim_u_spl'] + batch.INFLOWMETAINF['dim_CT'] + batch.INFLOWMETAINF['dim_NCC'] + batch.INFLOWMETAINF['dim_BatchEmb']
            ]
            ten_input_gnn = torch.cat(
                [ten_input_gnn, batch.y[:, rng_batchemb[0]:rng_batchemb[1]].to(device).detach()],
                1
            )

        gnn_output = self.module_gnn(
            ten_input_gnn,
            batch.edge_index.to(device)
        )

        if self.flag_enable_batchtoken:
            if self.param_batchshifts_4_gnnoutput is not None:  # it's the case when having one batch
                extended_batchemb = torch.cat(
                    [torch.zeros(self.module_gnn.dim_output, requires_grad=False).to(gnn_output.device).unsqueeze(0), self.param_batchshifts_4_gnnoutput],
                    0
                )  # [num_batch x dim_out]

                ten_shift = torch.mm(
                    batch.y[:, rng_batchemb[0]:rng_batchemb[1]].to(device).detach(),
                    extended_batchemb
                )  # [N x dim_out]
                gnn_output = gnn_output + ten_shift

        return gnn_output

    def _feed_to_head(self, str_int_or_spl, output_gnn, batch, device):
        '''
        Feeds to int or spl heads "not in the twosep mode".
        :param str_int_or_spl:
        :param output_gnn:
        :param batch:
        :param device:
        :return:
        '''
        assert (
            self.mode_headxint_headxspl_headboth_twosep != ModeArch.TWOSEP
        )
        assert (
            str_int_or_spl in ['int', 'spl']
        )
        module_mu = self.module_muxint if(str_int_or_spl == 'int') else self.module_muxspl
        module_cov = self.module_covxint if(str_int_or_spl == 'int') else self.module_covxspl

        # handle the None modules
        if module_mu is None:
            assert module_cov is None
            return None, None

        # create input_head
        input_head = self._extend_input(
            ten_input=output_gnn,
            key_inspoint=ArchInsertionPoint.HEADINT if(str_int_or_spl == 'int') else ArchInsertionPoint.HEADSPL,
            batch=batch,
            device=device
        )

        return module_mu(input_head), module_cov(input_head)


    def _feed_to_heads_twosep(self, x, str_int_or_spl, output_gnn, batch, device):
        '''
        Feeds to int or spl heads "in the twosep mode".
        :param x
        :param str_int_or_spl:
        :param output_gnn:
        :param batch:
        :param device:
        :return:
        '''
        assert (
            self.mode_headxint_headxspl_headboth_twosep == ModeArch.TWOSEP
        )
        assert (
            str_int_or_spl in ['int', 'spl']
        )
        module_mu = self.module_muxint if (str_int_or_spl == 'int') else self.module_muxspl
        module_cov = self.module_covxint if (str_int_or_spl == 'int') else self.module_covxspl
        assert module_mu is not None
        assert module_cov is not None

        input_head = self._extend_input(
            ten_input=x if(str_int_or_spl == 'int') else output_gnn,
            key_inspoint=ArchInsertionPoint.HEADINT if (str_int_or_spl == 'int') else ArchInsertionPoint.HEADSPL,
            batch=batch,
            device=device
        )

        return module_mu(input_head), module_cov(input_head)


    def _extend_input(self, ten_input:torch.Tensor, key_inspoint, batch, device):
        '''
        If CT or NCC are meant to be added to 'key_inspoint', they are concatenated to ten_input
        :param device:
        :param batch:
        :param ten_input:
        :param key_inspoint:
        :return:
        '''
        assert (isinstance(ten_input, torch.Tensor))
        assert (key_inspoint in [ArchInsertionPoint.NONE, ArchInsertionPoint.HEADINT, ArchInsertionPoint.HEADSPL, ArchInsertionPoint.BACKBONE])
        output = [ten_input]

        if self.dict_CTNNC_usage['CT'] == key_inspoint:
            rng_CT = [
                batch.INFLOWMETAINF['dim_u_int'] + batch.INFLOWMETAINF['dim_u_spl'],
                batch.INFLOWMETAINF['dim_u_int'] + batch.INFLOWMETAINF['dim_u_spl'] + batch.INFLOWMETAINF['dim_CT']
            ]
            output.append(
                batch.y[
                    :,
                    rng_CT[0]:rng_CT[1]
                ].to(device)
            )

        if self.dict_CTNNC_usage['NCC'] == key_inspoint:
            rng_NCC = [
                batch.INFLOWMETAINF['dim_u_int'] + batch.INFLOWMETAINF['dim_u_spl'] + batch.INFLOWMETAINF['dim_CT'],
                batch.INFLOWMETAINF['dim_u_int'] + batch.INFLOWMETAINF['dim_u_spl'] + batch.INFLOWMETAINF['dim_CT'] + batch.INFLOWMETAINF['dim_NCC']
            ]
            output.append(
                batch.y[
                    :,
                    rng_NCC[0]:rng_NCC[1]
                ].to(device)
            )

        output = torch.cat(output, 1) if (len(output) > 1) else output[0]
        return output

    def forward(self, batch, prob_maskknowngenes:float, ten_xy_absolute:torch.Tensor):
        '''
        :param batch:
        :param prob_maskknowngenes: must be zero for `Disentangler`, this arg is kept for consistency.
        :param ten_xy_absolute:
        :return:
        '''

        # make x_input
        with torch.no_grad():
            assert(prob_maskknowngenes == 0.0)
            x_log1p = torch.log(
                1.0 + batch.x.to_dense()  # TODO: how to make sure that batch.x contains the count data ???
            ).to(ten_xy_absolute.device)
            x_cnt = batch.x.to_dense().to(ten_xy_absolute.device).detach() + 0.0
            x_input = x_cnt if (self.str_mode_normalizex == 'counts') else x_log1p



        # feed to GNN
        output_gnn = self._feed_to_GNN(
            x_input=x_input,
            batch=batch,
            device=ten_xy_absolute.device
        )

        # the non-zero mask
        with torch.no_grad():
            oneon_x_nonzero = (x_cnt > 0.0) + 0  # [N, num_genes]

        # create output in modes other than TWOSEP ===
        if self.mode_headxint_headxspl_headboth_twosep != ModeArch.TWOSEP:
            EPS_COV = 1e-4 * torch.ones(
                size=[x_cnt.size()[0], x_cnt.size()[1]],
                device=ten_xy_absolute.device,
                requires_grad=False
            )

            muxint, covint = self._feed_to_head(
                str_int_or_spl='int',
                output_gnn=output_gnn,
                batch=batch,
                device=ten_xy_absolute.device
            )

            muxspl, covspl = self._feed_to_head(
                str_int_or_spl='spl',
                output_gnn=output_gnn,
                batch=batch,
                device=ten_xy_absolute.device
            )

            if muxint is not None:
                muxint = torch.clamp(
                    muxint * oneon_x_nonzero,
                    min=torch.tensor([0.0001], device=ten_xy_absolute.device),  # TODO: maybe tune?
                    max=x_cnt
                )  # [N, num_genes]

            if muxspl is not None:
                muxspl = torch.clamp(
                    muxspl * oneon_x_nonzero,
                    min=torch.tensor([0.0001], device=ten_xy_absolute.device),  # TODO: maybe tune?
                    max=x_cnt
                )  # [N, num_genes]

            if muxint is None:
                assert covint is None
                muxint = x_cnt - muxspl
                covint = covspl # EPS_COV #TODO: choose between covspl or EPS_COV?

            if muxspl is None:
                assert covspl is None
                muxspl = x_cnt - muxint
                covspl = covint # EPS_COV #TODO: choose between covspl or EPS_COV?
        else:
            assert self.mode_headxint_headxspl_headboth_twosep == ModeArch.TWOSEP
            muxint, covint = self._feed_to_heads_twosep(
                x=x_input,
                str_int_or_spl='int',
                output_gnn=output_gnn,
                batch=batch,
                device=ten_xy_absolute.device
            )
            muxspl, covspl = self._feed_to_heads_twosep(
                x=x_input,
                str_int_or_spl='spl',
                output_gnn=output_gnn,
                batch=batch,
                device=ten_xy_absolute.device
            )

            if muxint is not None:
                muxint = torch.clamp(
                    muxint * oneon_x_nonzero,
                    min=torch.tensor([0.0001], device=ten_xy_absolute.device),  # TODO: maybe tune?
                    max=x_cnt
                )  # [N, num_genes]

            if muxspl is not None:
                muxspl = torch.clamp(
                    muxspl * oneon_x_nonzero,
                    min=torch.tensor([0.0001], device=ten_xy_absolute.device),  # TODO: maybe tune?
                    max=x_cnt
                )  # [N, num_genes]




        # assert (not torch.any(ten_manually_masked))
        loss_imputex = None
        ten_out_imputer = 0.0



        # compute sigmaxint, sigmaxspl
        sigmaxint_raw = torch.clamp(
            torch.exp(
                covint
            ),
            min=0.0 * torch.ones_like(covint),  # TODO: maybe tune?
            max=(x_cnt ** 2).detach()
        )  # [N, num_genes]
        sigmaxspl_raw = torch.clamp(
            torch.exp(
                covspl
            ),
            min=0.0 * torch.ones_like(covspl),  # TODO: maybe tune?
            max=(x_cnt ** 2).detach()
        )  # [N, num_genes]

        '''
        The below part puts a lower bound on the variance on non-central nodes.
        This is done to solve the conceptual issue of getting GNN output on non-central nodes.
        '''

        sigmaxint = torch.clamp(sigmaxint_raw, min=self.std_minval_finalclip, max=self.std_maxval_finalclip)
        sigmaxspl = torch.clamp(sigmaxspl_raw, min=self.std_minval_finalclip, max=self.std_maxval_finalclip)

        sigmaxint = torch.concat(
            [
                sigmaxint[:,0:batch.batch_size]+0.0,
                torch.clamp(
                    sigmaxint[:,batch.batch_size::]+0.0,
                    min=self.clipval_cov_noncentralnodes,
                    max=4.0
                )],
            1
        ).sqrt()
        sigmaxspl = torch.concat(
            [
                sigmaxspl[:, 0:batch.batch_size]+0.0,
                torch.clamp(
                    sigmaxspl[:, batch.batch_size::]+0.0,
                    min=self.clipval_cov_noncentralnodes,
                    max=4.0
                )],
            1
        ).sqrt()



        # sigma-s cannot be more than the observed count
        #DONE above sigmaxint = torch.clamp(sigmaxint, max=x_cnt.detach())
        #DONE above sigmaxspl = torch.clamp(sigmaxspl, max=x_cnt.detach())


        return dict(
            x_cnt=x_cnt,
            muxint=muxint,
            muxspl=muxspl,
            sigmaxint=sigmaxint,
            sigmaxspl=sigmaxspl,
            ten_out_imputer=ten_out_imputer + 0.0,
            loss_imputex=loss_imputex
        )



