# from MLPerf training https://github.com/facebookresearch/dlrm/blob/6d75c84d834380a365e2f03d4838bee464157516/dlrm_s_pytorch.py

import contextlib
import itertools
import functools
import time
import json
from typing import Tuple
from collections import OrderedDict
# onnx
# The onnx import causes deprecation warnings every time workers
# are spawned during testing. So, we filter out those warnings.
import warnings
with warnings.catch_warnings():
    warnings.filterwarnings("ignore", category=DeprecationWarning)
import onnx

#import extend_distributed as ext_dist
from workloads.dlrm.workload import Dlrm

import spec 
import torch 
import torch.nn as nn
from torch.nn.parallel.parallel_apply import parallel_apply 
from torch.nn.parallel.scatter_gather import gather, scatter

import torch.nn.functional as F
import sklearn.metrics
from torch.optim.lr_scheduler import _LRScheduler



class _Dlrm_model(nn.Module):
    def create_mlp(self, ln, sigmoid_layer):
        # build MLP layer by layer
        layers = nn.ModuleList()
        for i in range(0, ln.size -1):
            n = ln[i]
            m = ln[i + 1]

            # construct fully connected operator 
            LL = nn.Linear(int(n), int(m), bias=True)
            
            # initialize the weights
            # with torch.no_grad():
            # custom Xavier input, output or two-sided fill
            mean = 0.0
            std_dev = np.sqrt(2/( m + n)) 
            W = np.random.normal(mean, std_dev, size=(m,n)).astype(np.float32)
            std_dev = np.sqrt(1 / m)
            bt = np.random.normal(mean, std_dev, size=m).astype(np.float32)

            LL.weight.data = torch.tensor(W, requires_grad=True)
            LL.bias.data = torch.tensor(bt, requires_grad=True)

            layers.append(LL)
            

            if i == sigmoid_layer:
                layers.append(nn.Sigmoid())
            else:
                layers.append(nn.ReLU())

        return torch.nn.Sequential(*layers)

    
    def create_emb(self, m, ln, weighted_pooling = None):
        emb_l = nn.ModuleList()
        for i in range(0, ln.size):
            n = ln[i]
            # construct embedding operator 
            if self.qr_flag and n > self.qr_threshold:
                EE = QREmbddingBag(n, 
                        m, 
                        self.qr_collisions, 
                        operation = self.qr_operation, 
                        mode = "sum", 
                        sparse = True)

            elif self.md_flag:
                base = max(m)
                _m = m[i] if n > self.md_threshold else base
                EE = PrEmbeddingBag(n, _m, base)
                # use numpy for consistency 
                W = np.random.uniform(
                        low = -np.sqrt(1/n),
                        high = np.sqrt(1/n),
                        size =(m,n)).astype(np.float32)

                EE.embs.weight.data = torch.tensor(W, requires_grad=True)
            else:
                EE = nn.EmbeddingBag(n, m, mode="sum", sparse = True)

                #initialize embedding 
                W = np.random.uniform(
                        low = -np.sqrt(1/n),
                        high = np.sqrt(1/n),
                        size =(m,n)).astype(np.float32)

                EE.weight.data = torch.tensor(W, requires_grad=True)
            emb_l.append(EE)
        return emb_l

    def __init__(
        self,
        m_spa = None,
        ln_emb = None,
        ln_bot = None,
        ln_top = None,
        arch_interaction_op = None,
        arch_interaction_itself = False,
        sigmoid_bot = -1,
        sigmoid_top=-1,
        sync_dense_params = True,
        loss_threshold = 0.0,
        ndevices=-1,
        qr_flag = False,
        qr_operation = "mult",
        qr_collisions = 0,
        qr_threshold = 200,
        md_flag = False,
        md_threshold = 200
    ):
        super(_Dlrm_model, self).__init__()
        if (
            (m_spa is not None)
            and (ln_emb is not None)
            and (ln_bot is not None)
            and (ln_top is not None)
            and (arch_interaction_op is not None)
        ):
            self.ndevices = ndevices
            self.output_d = 0
            self.parallel_model_batch_size = -1
            self.parallel_model_is_not_prepared = True
            self.arch_interaction_op = arch_interaction_op
            self.arch_interaction_itself = arch_interaction_itself
            self.sync_dense_params = sync_dense_params
            self.loss_threshold = loss_threshold
            
            # create variables for QR embedding if applicable
            self.qr_flag = qr_flag
            if self.qr_flag:
                self.qr_collisions = qr_collisions
                self.qr_operation = qr_operation
                self.qr_threshold = rq_threshold

            # create variables for MD embedding if applicable
            if ndevices <= 1:
                self.emb_l = self.create_emb(m_spa, ln_emb)
            self.bot_l = self.create_mlp(ln_bot, sigmoid_bot)
            self.top_l = self.create_mlp(ln_top, sigmoid_top)




    def apply_mlp(self, x, layers):
        return layers(x)



    def apply_emb(self, lS_o, lS_i, emb_l):
        # WARNING: notice that we are processing the batch at once. We implicitly
        # assume that the data is laid out such that:
        # 1. each embedding is indexed with a group of sparse indices,
        #   corresponding to a single lookup
        # 2. for each embedding the lookups are further organized into a batch
        # 3. for a list of embedding tables there is a list of batched lookups

        ly = []
        for k, sparse_index_group_batch in enumerate(lS_i):
            sparse_offset_group_batch = lS_o[k]

            E = emb_l[k]
            V = E(sparse_index_group_batch, sparse_offset_group_batch)

        
            ly.append(V)

        #print(ly)
        return ly
    

    def interact_features(self, x, ly):
        if self.arch_interaction_op == "dot":
            # concatenate dense and sparse features
            (batch_size, d) = x.shape
            T = torch.cat([x] + ly, dim=1).view((batch_size, -1, d))
            # perform dot product 
            Z = torch.bmm(T, torch.transpose(T, 1, 2))
            # append dense feature with the interactions (into a row vector)
            _, ni, nj = Z.shape
            
            offset = 1 if self.atch_interaction_itself else 0
            li = torch.tensor([i for i in range(ni) for j in range(i + offset)])
            lj = torch.tensor([j for i in range(nj) for j in range(i + offset)])
            Zflat = Z[:, li, lj]

            R = torch.cat([x] + [Zflat], dim =1)
        elif self.arch_interaction_op == "cat":
            # concatenation features (into a row vector)
            R = torch.cat([x] + [Zflat], dim =1)
        else:
             sys.exit(
                "ERROR: --arch-interaction-op="
                + self.arch_interaction_op
                + " is not supported"
            )


        return R

    def sequential_forward(self, dense_x, lS_o, lS_i):
        # process dense features (using bottom mlp), resulting in a row vector
        x = self.apply_mlp(dense_x, self.bot_l)
        # debug prints
        # print("intermediate")
        # print(x.detach().cpu().numpy())

        # process sparse features(using embeddings), resulting in a list of row vectors
        ly = self.apply_emb(lS_o, ls_i, self.emb_l)

        # interact features (dense and sparse)
        z = self.interact_features(x, ly)
        # print(z.detach().cpu().numpy())

        # obtain probability of click (using top mlp)
        p = self.apply_mlp(z, self.top_l)

        # clamp output if needed
        if 0.0 < self.loss_threshold and self.loss_threshold < 1.0:
            z = torch.clamp(p, min=self.loss_threshold, max=(1.0 - self.loss_threshold)) 
        else:
            z = p

        return z








class DlrmWorkload(Dlrm):
    def __init__(self):
        self._eval_ds = None


    def build_input_queue(
        self, 
        data_rng: spec.RandomState,
        split: str,
        data_dir: str,
        batch_size: int):
        #return iter(self._build_dataset(data_rng, split, data_dir, batch_size))
        pass

    
    def eval_model(
        self,
        params: spec.ParameterContainer,
        model_state: spec.ModelAuxiliaryState,
        rng: spec.RandomState,
        data_dir: str):
        """Run a full evaluation of the model."""
        pass

   
    @property
    def eval_period_time_sec(self):
        pass


    def init_model_fn(
        self,
        rng: spec.RandomState):
        pass

    def is_output_params(self, param_key: spec.ParameterKey) -> bool:
        pass

    # Does NOT apply regularization, which is left to the submitter to do in
    # `update_params`.
    def loss_fn(
        self,
        label_batch: spec.Tensor,
        logits_batch: spec.Tensor) -> spec.Tensor: # differentiable
        pass

    @property  
    def max_allowed_runtime_sec(self) -> int: 
        pass
    
    def model_fn(
        self,
        params: spec.ParameterContainer,
        input_batch: spec.Tensor,
        model_state: spec.ModelAuxiliaryState,
        mode: spec.ForwardPassMode,
        rng: spec.RandomState,
        update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
        pass

    def model_params_types(self):
        pass

    # Keep this separate from the loss function in order to support optimizers
    # that use the logits.
    def output_activation_fn(
        self,
        logits_batch: spec.Tensor,
        loss_type: spec.LossType) -> spec.Tensor:
        """Return the final activations of the model."""
        pass
    
    @property 
    def param_shapes(self):
        pass
    
    @property
    def train_mean(self):
        pass

    @property
    def train_stddev(self):
        pass


    @property
    def num_eval_examples(self):
       """The size of the evaluation set."""
       pass
 
    @property
    def num_train_examples(self):
        """The size of the training set."""
        pass
   
    @property
    def target_value(self):
        """The target value to reach."""
        pass

