# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from __future__ import division
from __future__ import print_function

import numpy as np
import pandas as pd
import copy
import random
from .pytorch_utils import count_parameters
from ...utils import (
    get_or_create_path,
)
from ...log import get_module_logger

import torch
import torch.nn as nn
import torch.optim as optim

from ...model.base import Model
from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP


class IPPE(Model):
    """IPPE Model

    Parameters
    ----------
    d_feat : int
        input dimension for each time step
    metric: str
        the evaluate metric used in early stop
    optimizer : str
        optimizer name
    GPU : str
        the GPU ID(s) used for training
    """

    def __init__(
        self,
        d_feat=6,
        hidden_size=64,
        num_layers=2,
        dropout=0.0,
        n_epochs=200,
        batch_size=100,
        early_stop=5,
        loss="mse",
        output_dim=5,
        lr=5e-4,
        use_pretrain_pair=True,
        use_pretrain_point=True,
        pretrain_epoch=20,
        pretrain_lr=2e-4,
        point_enhance=True,
        pair_enhance=True,
        point_dim=10,
        pair_dim=10,
        iter_num=1,
        evaluate_time=3,
        inference_time=10,
        metric="loss",
        optimizer="adam",
        GPU=0,
        seed=None,
        **kwargs
    ):
        # Set logger.
        self.logger = get_module_logger("IPPE")
        self.logger.info("IPPE pytorch version...")

        # set hyper-parameters.
        self.d_feat = d_feat
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.n_epochs = n_epochs
        self.batch_size = batch_size
        self.early_stop = early_stop
        self.loss = loss
        self.output_dim = output_dim
        self.lr = lr
        self.use_pretrain_pair = use_pretrain_pair
        self.use_pretrain_point = use_pretrain_point
        self.pretrain_epoch = pretrain_epoch
        self.pretrain_lr = pretrain_lr
        self.point_enhance = point_enhance
        self.pair_enhance = pair_enhance
        self.point_dim = point_dim
        self.pair_dim = pair_dim
        self.iter_num = iter_num
        self.evaluate_time = evaluate_time
        self.inference_time = inference_time
        self.optimizer = optimizer
        self.metric = metric
        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu")
        self.seed = seed

        self.logger.info(
            "IPPE parameters setting:"
            "\nd_feat : {}"
            "\nhidden_size : {}"
            "\nnum_layers : {}"
            "\ndropout : {}"
            "\nn_epochs : {}"
            "\nbatch_size : {}"
            "\nearly_stop : {}"
            "\nloss_type : {}"
            "\noutput_dim : {}"
            "\nlr : {}"
            "\nuse_pretrain_pair : {}"
            "\nuse_pretrain_point : {}"
            "\npretrain_epoch : {}"
            "\npretrain_lr : {}"
            "\npoint_enhance : {}"
            "\npair_enhance : {}"
            "\npoint_dim : {}"
            "\npair_dim : {}"
            "\niter_num : {}"
            "\ninference_time : {}"
            "\noptimizer : {}"
            "\nseed : {}".format(
                d_feat,
                hidden_size,
                num_layers,
                dropout,
                n_epochs,
                batch_size,
                early_stop,
                loss,
                output_dim,
                lr,
                use_pretrain_pair,
                use_pretrain_point,
                pretrain_epoch,
                pretrain_lr,
                point_enhance,
                pair_enhance,
                point_dim,
                pair_dim,
                iter_num,
                inference_time,
                optimizer,
                seed,
            )
        )

        if self.seed is not None:
            np.random.seed(self.seed)
            torch.manual_seed(self.seed)
        self.point_model = GRU_POINT(
            d_feat=self.d_feat,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
        ).to(self.device)
        self.pair_model = GRU_PAIR(
            d_feat=self.d_feat*2,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
        ).to(self.device)
        self.model = MUTUAL_ENHANCE(device=self.device,
                                    point_input=self.hidden_size, \
                                    pair_input=self.hidden_size, \
                                    point_dim=self.point_dim, \
                                    pair_dim=self.pair_dim, \
                                    point_enhance=self.point_enhance, \
                                    pair_enhance=self.pair_enhance).to(self.device)
        for param in self.model.parameters():
            param.requires_grad = True

        param_list = list(self.model.parameters()) + list(self.point_model.parameters()) + list(self.pair_model.parameters())

        self.logger.info("IPPE model size: {:.4f} MB".format(count_parameters(self.model)))

        if optimizer.lower() == "adam":
            self.train_optimizer = optim.Adam(param_list, lr=self.lr)
            if self.use_pretrain_point:
                self.point_pretrain_optimizer = optim.Adam(self.point_model.parameters(), lr=self.pretrain_lr)
            if self.use_pretrain_pair:
                self.pair_pretrain_optimizer = optim.Adam(self.pair_model.parameters(), lr=self.pretrain_lr)
        elif optimizer.lower() == "gd":
            self.train_optimizer = optim.SGD(param_list, lr=self.lr)
            if self.use_pretrain_point:
                self.point_pretrain_optimizer = optim.SGD(self.point_model.parameters(), lr=self.pretrain_lr)
            if self.use_pretrain_pair:
                self.pair_pretrain_optimizer = optim.SGD(self.pair_model.parameters(), lr=self.pretrain_lr)
        else:
            raise NotImplementedError("optimizer {} is not supported!".format(optimizer))

        self.fitted = False

    @property
    def use_gpu(self):
        return self.device != torch.device("cpu")

    def mse(self, pred, label):
        loss = (pred - label) ** 2
        return torch.mean(loss)

    def loss_fn(self, pred, label):
        mask = ~torch.isnan(label)

        if self.loss == "mse":
            return self.mse(pred[mask], label[mask])

        raise ValueError("unknown loss `%s`" % self.loss)

    def metric_fn(self, pred, label):

        mask = torch.isfinite(label)

        if self.metric == "" or self.metric == "loss":
            return -self.loss_fn(pred[mask], label[mask])

        raise ValueError("unknown metric `%s`" % self.metric)

    def pair_data_process(self, feature, label):
        
        feature = np.array(feature)
        label = np.array(label)
        stock_num = feature.shape[0]

        pair_label = []
        if label.shape[0]:
            vec_1 = np.repeat(label, stock_num).reshape(stock_num, stock_num)
            vec_2 = np.transpose(vec_1)
            pair_label = (vec_1 - vec_2).reshape(-1)

        loc_1 = np.repeat(range(stock_num), stock_num).reshape(stock_num, stock_num)
        loc_2 = np.transpose(loc_1)
        loc_1 = loc_1.reshape([-1])
        loc_2 = loc_2.reshape([-1])

        return feature[loc_1], feature[loc_2], pair_label


    def train_epoch(self, x_train, y_train):

        self.model.train(True)
        self.point_model.train(True)
        self.pair_model.train(True)
        
        days = x_train.index.get_level_values(0).unique()
        total_day = days.shape[0]
        day_index_seq = np.array(range(total_day))
        random.shuffle(day_index_seq)

        for _, day_index in enumerate(day_index_seq):
            # prepare feature and label
            day = days[day_index]
            feature = np.array(x_train.loc[day])
            label = np.array(y_train.loc[day]).reshape([-1])
            stock_num = feature.shape[0]

            loc = np.random.choice(stock_num, self.batch_size)
            feature = feature[loc]
            label = label[loc]

            pair_feature1, pair_feature2, pair_label = self.pair_data_process(feature, label)
            
            point_feature = torch.tensor(feature, dtype=torch.float, device=self.device)
            point_label = torch.tensor(label, dtype=torch.float, device=self.device)
            pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
            pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)
            pair_label = torch.tensor(pair_label, dtype=torch.float, device=self.device)

            # transfer feature to embedding
            point = self.point_model(point_feature)
            pair = self.pair_model(pair_feature1, pair_feature2)
            point = self.model.transform_point(point)
            pair = self.model.transform_pair(pair)

            point_loss = 0
            pair_loss = 0
            if self.point_enhance:
                point_pred = self.model.point2pred(point)
                point_loss += self.loss_fn(point_pred.view(-1), point_label.view(-1))
            if self.pair_enhance:
                pair_pred = self.model.pair2pred(pair).view(-1)
                pair_loss += self.loss_fn(pair_pred.view(-1), pair_label.view(-1))

            for _ in range(self.iter_num):
                point, pair = self.model(point, pair)

                if self.point_enhance:
                    point_pred = self.model.point2pred(point)
                    point_loss += self.loss_fn(point_pred.view(-1), point_label.view(-1))
                if self.pair_enhance:
                    pair_pred = self.model.pair2pred(pair).view(-1)
                    pair_loss += self.loss_fn(pair_pred.view(-1), pair_label.view(-1))

            loss = (point_loss + pair_loss)/((1 + self.iter_num)*self.iter_num/2)

            self.train_optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_value_(self.point_model.parameters(), 3.0)
            self.train_optimizer.step()

    def test_epoch(self, data_x, data_y):
        self.model.eval()
        self.point_model.eval()
        self.pair_model.eval()

        scores = []
        losses = []

        days = data_x.index.get_level_values(0).unique()
        total_day = days.shape[0]
        day_index_seq = np.array(range(total_day))

        for _ in range(self.evaluate_time):
            for _, day_index in enumerate(day_index_seq):
                # prepare feature and label
                day = days[day_index]
                feature = np.array(data_x.loc[day])
                label = np.array(data_y.loc[day]).reshape([-1])
                stock_num = feature.shape[0]

                loc = np.random.choice(stock_num, self.batch_size)
                feature = feature[loc]
                label = label[loc]

                pair_feature1, pair_feature2, pair_label = self.pair_data_process(feature, label)
                point_feature = torch.tensor(feature, dtype=torch.float, device=self.device)
                point_label = torch.tensor(label, dtype=torch.float, device=self.device)
                pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
                pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)
                pair_label = torch.tensor(pair_label, dtype=torch.float, device=self.device)

                with torch.no_grad():
                    # transfer feature to embedding
                    point = self.point_model(point_feature)
                    pair = self.pair_model(pair_feature1, pair_feature2)
                    point = self.model.transform_point(point)
                    pair = self.model.transform_pair(pair)

                    for _ in range(self.iter_num):
                        point, pair = self.model(point, pair)
                    point_pred = self.model.point2pred(point)

                    loss = self.loss_fn(point_pred, point_label)
                    losses.append(loss.item())
                    score = self.metric_fn(point_pred, point_label)
                    scores.append(score.item())

        return np.mean(losses), np.mean(scores)

    def pretrain_point_model(
        self,
        x_train,
        y_train,
        batch_size=2000,
    ):

        x_train_values = x_train.values
        y_train_values = np.squeeze(y_train.values)

        self.point_model.train()

        for step in range(self.pretrain_epoch):
            self.logger.info("Epoch%d:", step)
            indices = np.arange(len(x_train_values))
            np.random.shuffle(indices)

            for i in range(len(indices))[:: batch_size]:

                if len(indices) - i < batch_size:
                    break

                feature = torch.from_numpy(x_train_values[indices[i : i + batch_size]]).float().to(self.device)
                label = torch.from_numpy(y_train_values[indices[i : i + batch_size]]).float().to(self.device)

                point_embedding = self.point_model(feature)
                pred = self.point_model.fc_out(point_embedding).squeeze()
                loss = self.loss_fn(pred, label)

                self.point_pretrain_optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_value_(self.point_model.parameters(), 3.0)
                self.point_pretrain_optimizer.step()



    def pretrain_pair_model(
        self,
        x_train,
        y_train,
    ):

        days = x_train.index.get_level_values(0).unique()
        total_day = days.shape[0]
        day_index_seq = np.array(range(total_day))
        random.shuffle(day_index_seq)

        for step in range(self.pretrain_epoch):
            self.logger.info("Epoch%d:", step)
            for _, day_index in enumerate(day_index_seq):
                # prepare feature and label
                day = days[day_index]
                feature = np.array(x_train.loc[day])
                label = np.array(y_train.loc[day]).reshape([-1])
                stock_num = feature.shape[0]

                loc = np.random.choice(stock_num, self.batch_size)
                feature = feature[loc]
                label = label[loc]

                pair_feature1, pair_feature2, pair_label = self.pair_data_process(feature, label)
                
                pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
                pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)
                pair_label = torch.tensor(pair_label, dtype=torch.float, device=self.device)

                pair_embedding = self.pair_model(pair_feature1, pair_feature2)
                pair_pred = self.pair_model.fc_out(pair_embedding).squeeze()
                loss = self.loss_fn(pair_pred, pair_label)

                self.pair_pretrain_optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_value_(self.pair_model.parameters(), 3.0)
                self.pair_pretrain_optimizer.step()

    def fit(
        self,
        dataset: DatasetH,
        evals_result=dict(),
        save_path=None,
    ):

        df_train, df_valid, _ = dataset.prepare(
            ["train", "valid", "test"],
            col_set=["feature", "label"],
            data_key=DataHandlerLP.DK_L,
        )
        save_path = get_or_create_path(save_path)
        stop_steps = 0
        best_score = -np.inf
        best_epoch = 0
        evals_result["valid"] = []

        self.logger.info("training...")
        self.fitted = True

        if self.use_pretrain_point:
            self.logger.info("pretrain point model...")
            self.pretrain_point_model(df_train["feature"], df_train["label"])

        if self.use_pretrain_pair:
            self.logger.info("pretrain pair model...")
            self.pretrain_pair_model(df_train["feature"], df_train["label"])

        for step in range(self.n_epochs):
            self.logger.info("Epoch%d:", step)
            self.train_epoch(df_train["feature"], df_train["label"])
            _, val_score = self.test_epoch(df_valid["feature"], df_valid["label"])
            self.logger.info("valid %.6f" % (val_score))
            evals_result["valid"].append(val_score)

            if val_score > best_score:
                best_score = val_score
                stop_steps = 0
                best_epoch = step
                best_param = [copy.deepcopy(self.model.state_dict()),\
                                copy.deepcopy(self.point_model.state_dict()),\
                                copy.deepcopy(self.pair_model.state_dict())]

            else:
                stop_steps += 1
                if stop_steps >= self.early_stop:
                    self.logger.info("early stop")
                    break

        self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch))
        self.model.load_state_dict(best_param[0])
        self.point_model.load_state_dict(best_param[1])
        self.pair_model.load_state_dict(best_param[2])
        torch.save(best_param[0], save_path+'ippe')
        torch.save(best_param[1], save_path+'point_model')
        torch.save(best_param[2], save_path+'pair_model')

    def predict(self, dataset: DatasetH):
        if not self.fitted:
            raise ValueError("model is not fitted yet!")

        test_data = dataset.prepare("test", col_set=["feature", "label"], data_key=DataHandlerLP.DK_I)
        x_test = test_data["feature"]
        y_test = test_data["label"]
        test_data.to_pickle("~/tmp_test_data.pkl")
        self.model.eval()
        self.point_model.eval()
        self.pair_model.eval()

        preds = []
        indexes = []
        days = x_test.index.get_level_values(0).unique()
        total_day = days.shape[0]

        for day_index in range(total_day):
            # prepare feature and label
            day = days[day_index]
            daily_feature = np.array(x_test.loc[day])
            daily_index = x_test.loc(axis=0)[day,:].index
            stock_num = daily_feature.shape[0]

            for _ in range(self.inference_time):
                loc = np.random.choice(stock_num, self.batch_size)
                feature = daily_feature[loc]
                currend_index = daily_index[loc]

                pair_feature1, pair_feature2, _ = self.pair_data_process(feature, [])
                point_feature = torch.tensor(feature, dtype=torch.float, device=self.device)
                pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
                pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)

                with torch.no_grad():
                    # transfer feature to embedding
                    point = self.point_model(point_feature)
                    pair = self.pair_model(pair_feature1, pair_feature2)
                    point = self.model.transform_point(point)
                    pair = self.model.transform_pair(pair)

                    for _ in range(self.iter_num):
                        point, pair = self.model(point, pair)
                    point_pred = self.model.point2pred(point)

                preds.append(point_pred.detach().cpu().numpy().reshape([-1]))
                indexes.append(currend_index)

        res_df = pd.Series(np.concatenate(preds), index=np.concatenate(indexes))
        res_df.index = pd.MultiIndex.from_tuples(res_df.index)
        res_df.index.names = ['datetime', 'instrument']
        res_df = res_df.groupby(level=[0,1]).apply(lambda x:x.mean())
        preds = pd.Series(np.zeros(x_test.shape[0]), index=test_data.index)
        preds.loc[res_df.index] = np.array(res_df)

        return preds


class GRU_POINT(nn.Module):
    
    def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0):
        super().__init__()

        self.rnn = nn.GRU(
            input_size=d_feat,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
        )
        self.fc_out = nn.Linear(hidden_size, 1)

        self.d_feat = d_feat

    def forward(self, x):
        # x: [N, F*T]
        x = x.reshape(len(x), self.d_feat, -1) # [N, F, T]
        x = x.permute(0, 2, 1) # [N, T, F]
        out, _ = self.rnn(x)
        return out[:, -1, :].squeeze()

class GRU_PAIR(nn.Module):

    def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0):
        super().__init__()

        self.rnn = nn.GRU(
            input_size=d_feat*2,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
        )
        self.fc_out = nn.Linear(hidden_size, 1)
        self.d_feat = d_feat

    def forward(self, x, y):
        # x: [N, F*T]
        x = x.reshape(len(x), self.d_feat, -1) # [N, F, T]
        x = x.permute(0, 2, 1) # [N, T, F]
        y = y.reshape(len(y), self.d_feat, -1) # [N, F, T]
        y = y.permute(0, 2, 1) # [N, T, F]
        
        input = torch.cat((x, y), 2)
        out, _ = self.rnn(input)
        return out[:, -1, :].squeeze()

class MUTUAL_ENHANCE(nn.Module):
    def __init__(self, device, point_input, pair_input, point_dim=10, pair_dim=10, point_enhance=True, pair_enhance=True):
        
        super().__init__()
        self.device = device
        self.point_dim = point_dim
        self.pair_dim = pair_dim
        self.point_enhance = point_enhance
        self.pair_enhance = pair_enhance
        self.transform_point = nn.Linear(point_input, point_dim)
        self.transform_pair = nn.Linear(pair_input, pair_dim)
        self.point2pred = nn.Linear(point_dim, 1)
        self.pair2pred = nn.Linear(pair_dim, 1)

        # prepare lamda block
        self.transferion_point = torch.nn.Parameter(torch.rand(self.point_dim, self.point_dim*self.pair_dim).cuda(), \
                                             requires_grad=True)
        
        self.transferion_pair = torch.nn.Parameter(torch.rand(self.point_dim, self.pair_dim*self.point_dim).cuda(), \
                                             requires_grad=True)

        self.lamda_point_func = nn.Sequential(
            nn.Linear(3*self.point_dim+3, 16),
            nn.ReLU(),
            nn.Linear(16,16),
            nn.ReLU(),
            nn.Linear(16,1),
            nn.Sigmoid()) 

        self.omega_point_func = nn.Sequential(
            nn.Linear(3*self.point_dim + 2*self.pair_dim + 2, 32),
            nn.ReLU(),
            nn.Linear(32,32),
            nn.ReLU(),
            nn.Linear(32,16),
            nn.ReLU(),
            nn.Linear(16,1))

        self.lamda_pair_func = nn.Sequential(
            nn.Linear(2*self.point_dim + 3*self.pair_dim + 2,16),
            nn.ReLU(),
            nn.Linear(16,16),
            nn.ReLU(),
            nn.Linear(16,1),
            nn.Sigmoid()) 

        self.softmax = torch.nn.Softmax(dim = 1)
            
    def forward(self, point, pair):
        # data prepare

        stock_num = point.shape[0]
        point_1 = point.expand(stock_num, stock_num, point.shape[1])
        point_2 = torch.transpose(point_1,0,1)
        pair = pair.view(stock_num, stock_num, -1)
              
        pair_1 = pair
        pair_2 = torch.transpose(pair_1,0,1)
        
        if self.point_enhance:
            # calculate point estimation embedding
            correction_in_1 = pair_1.view(stock_num, stock_num, self.pair_dim, 1)
            correction_in_2 = point_1.view(stock_num, stock_num, 1, self.point_dim)
            tmp = torch.matmul(correction_in_2, self.transferion_point)
            tmp = tmp.view(stock_num, stock_num, self.point_dim, self.pair_dim)
            self.point_estimation = torch.matmul(tmp, correction_in_1).view(stock_num, stock_num, self.point_dim)
            self.point_estimation_pred = self.point2pred(self.point_estimation)

            # omega in
            self.point_estimation_mean = torch.sum(self.point_estimation, dim=1)/stock_num
            self.point_estimation_var = torch.std(self.point_estimation, dim=1)
            self.point_estimation_pred_mean = torch.sum(self.point_estimation_pred, dim=1)/stock_num
            self.point_estimation_pred_var = torch.std(self.point_estimation_pred, dim=1)
            
            self.omega_in = torch.cat((pair_1, # embedding(i,j)
                                        pair_2, # embedding(j,i)
                                        point_1, # embedding(j)
                                        point_2, # embedding(i)
                                        self.point_estimation,
                                        self.point_estimation_pred_mean.expand(stock_num, stock_num).view(stock_num, stock_num, -1),
                                        self.point_estimation_pred_var.expand(stock_num, stock_num).view(stock_num, stock_num, -1),
                                  ), 2)

            # omega out
            self.omega_point = self.omega_point_func(self.omega_in)
            self.omega_point = self.softmax(self.omega_point)

            # calculate aggerated esetimation embedding
            self.agg_point_estimation = torch.sum(self.omega_point*self.point_estimation, dim=1)
            # self.agg_point_estimation = torch.sum(self.point_estimation, dim=1)/stock_num # average

            # lambda in
            agg_point_estimation_pred = self.point_estimation_pred*self.omega_point
            agg_point_estimation_pred_mean = torch.sum(agg_point_estimation_pred, dim=1)/stock_num
            agg_point_estimation_pred_var = torch.std(agg_point_estimation_pred, dim=1)
            omega_std = torch.std(self.omega_point, dim=1)
            distance = torch.abs(self.agg_point_estimation - point)
            self.lamda_in_point = torch.cat((point,
                              self.agg_point_estimation.view(stock_num, -1),
                              agg_point_estimation_pred_mean.view(stock_num, -1),
                              agg_point_estimation_pred_var.view(stock_num, -1),
                              omega_std.view(stock_num, -1),
                              distance.view(stock_num, -1)
                             ), 1)
            # lambda out
            self.lamda_point = self.lamda_point_func(self.lamda_in_point)
            # update
            point = (1-self.lamda_point)*point + self.lamda_point*self.agg_point_estimation

        if self.pair_enhance:
            # calculate pair estimation embedding
            correction_in_1 = point_1.view(stock_num, stock_num, 1, self.point_dim)
            correction_in_2 = point_2.view(stock_num, stock_num, self.point_dim, 1)
            tmp = torch.matmul(correction_in_1, self.transferion_pair)
            tmp = tmp.view(stock_num, stock_num, self.pair_dim, self.point_dim)
            self.pair_estimation = torch.matmul(tmp, correction_in_2).view(stock_num, stock_num, self.pair_dim)
            self.pair_estimation_pred = self.point2pred(self.pair_estimation)
        
            # lambda in
            self.pair_estimation_mean = torch.sum(self.pair_estimation, dim=1)/stock_num
            self.pair_estimation_var = torch.std(self.pair_estimation, dim=1)
            self.pair_estimation_pred_mean = torch.sum(self.pair_estimation_pred, dim=1)/stock_num
            self.pair_estimation_pred_var = torch.std(self.pair_estimation_pred, dim=1)
            
            self.lamda_in_pair = torch.cat((pair_1, # embedding(i,j)
                                        pair_2, # embedding(j,i)
                                        point_1, # embedding(j)
                                        point_2, # embedding(i)
                                        self.pair_estimation,
                                        self.pair_estimation_pred_mean.expand([stock_num, stock_num]).view(stock_num, stock_num, -1),
                                        self.pair_estimation_pred_var.expand([stock_num, stock_num]).view(stock_num, stock_num, -1),
                                  ), 2)

            # lambda out
            self.lambda_pair = self.lamda_pair_func(self.lamda_in_pair)
            # update
            pair = (1-self.lambda_pair)*pair + self.lambda_pair*self.pair_estimation

        return point, pair
