from flask import Blueprint, request
from big_data.db import get_db, close_db
import json
from big_data.spark import get_spark, get_client
import urllib.parse

from pyspark.ml.feature import StringIndexer
from pyspark.ml.feature import RFormula
from pyspark.ml.feature import StandardScaler

from pyspark.ml.classification import LogisticRegression
from pyspark.ml.classification import DecisionTreeClassifier
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.classification import GBTClassifier
from pyspark.ml.classification import NaiveBayes
from pyspark.ml.classification import LinearSVC
from pyspark.ml.regression import LinearRegression
from pyspark.ml.regression import RandomForestRegressor
from pyspark.ml.regression import GBTRegressor
from pyspark.ml.evaluation import BinaryClassificationEvaluator
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
from pyspark.ml.evaluation import RegressionEvaluator

import logging
LOG_FORMAT = "/********** %(asctime)s - %(levelname)s - %(message)s **********/"
logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)

bp = Blueprint('add-model', __name__, url_prefix='/add-model')

# 根据 dataset_id 获取 dataframe
def getDF(dataset_id):
    db = get_db()
    c = db.cursor()

    # 通过 dataset_id 获取 project_id 和 filename
    t = (dataset_id,)
    c.execute("SELECT project_id, filename FROM dataset WHERE id=?", t)
    res = c.fetchone()
    project_id = tuple(res)[0]
    filename = tuple(res)[1]

    # 通过 project_id 获取 username
    t = (project_id,)
    c.execute("SELECT username FROM project WHERE id=?", t)
    res = c.fetchone()
    username = tuple(res)[0]

    # 拼接出 hdfs_path，获得数据集 df
    hdfs_path = "/user/" + username + "/" + str(project_id) + "/" + filename + ".parquet"
    df = get_spark().read.load(hdfs_path)

    return df

# 数据预处理
def dataPreProcess(df, label, features, problem_type):
    # 把 df 中除 label 列的数据类型是字符串的各列先 StringIndexer 一下
    for i, dtype in enumerate(df.dtypes):
        if dtype[1] == "string" and i != int(label):
            df = df.withColumnRenamed(dtype[0], dtype[0] + "_old")
            indexer = StringIndexer(inputCol = dtype[0] + "_old", outputCol = dtype[0])
            df = indexer.fit(df).transform(df)

    # 如果是分类问题，用 StringIndexer 对 label 列进行处理
    if problem_type != 3:
        label_name = df.dtypes[int(label)][0]
        df = df.withColumnRenamed(label_name, label_name + "_old")
        indexer = StringIndexer(inputCol = label_name + "_old", outputCol = label_name)
        df = indexer.fit(df).transform(df)

    # 装配 label 和 features vector
    formula_str = df.columns[-1] + "~"
    for feature in features:
        formula_str = formula_str + df.columns[int(feature)] + "+"
    formula_str = formula_str[:-1]
    formula = RFormula(
        formula = formula_str,
        featuresCol = "features",
        labelCol = "label"
    )
    df = formula.fit(df).transform(df)

    # 特征归一化处理
    df = df.withColumnRenamed("features", "features_old")
    standardscaler = StandardScaler().setInputCol("features_old").setOutputCol("features")
    df = standardscaler.fit(df).transform(df)

    return df

# 根据参数配置分类器
def getClassifier(type, params):
    if type == 1: # Logistic Regression
        lr_aggregation_depth = int(params["lr_aggregation_depth"])
        lr_elastic_net_param = float(params["lr_elastic_net_param"])
        lr_fit_intercept = True if params["lr_fit_intercept"] == "true" else False
        lr_max_iter = int(params["lr_max_iter"])
        lr_reg_param = float(params["lr_reg_param"])
        cf = LogisticRegression(
            aggregationDepth = lr_aggregation_depth,
            elasticNetParam = lr_elastic_net_param,
            fitIntercept = lr_fit_intercept,
            maxIter = lr_max_iter,
            regParam = lr_reg_param
        )
    elif type == 2: # Decision Tree Classifier
        dt_max_depth = int(params["dt_max_depth"])
        dt_impurity = params["dt_impurity"]
        cf = DecisionTreeClassifier(
            maxDepth = dt_max_depth,
            impurity = dt_impurity
        )
    elif type == 3: # Random Forest Classifier
        rf_impurity = params["rf_impurity"]
        rf_max_depth = int(params["rf_max_depth"])
        rf_num_trees = int(params["rf_num_trees"])
        cf = RandomForestClassifier(
            impurity = rf_impurity,
            maxDepth = rf_max_depth,
            numTrees = rf_num_trees
        )
    elif type == 4: # Gradient-boosted Tree Classifier
        gbt_max_depth = int(params["gbt_max_depth"])
        gbt_max_iter = int(params["gbt_max_iter"])
        gbt_step_size = float(params["gbt_step_size"])
        gbt_subsampling_rate = float(params["gbt_subsampling_rate"])
        cf = GBTClassifier(
            maxDepth = gbt_max_depth,
            maxIter = gbt_max_iter,
            stepSize = gbt_step_size,
            subsamplingRate = gbt_subsampling_rate
        )
    elif type == 5: # Naive Bayes
        nb_smoothing = float(params["nb_smoothing"])
        nb_model_type = params["nb_model_type"]
        cf = NaiveBayes(
            smoothing = nb_smoothing,
            modelType = nb_model_type
        )
    elif type == 6: # Linear Support Vector Machine
        lsvc_max_iter = int(params["lsvc_max_iter"])
        lsvc_reg_param = float(params["lsvc_reg_param"])
        lsvc_fit_intercept = True if params["lsvc_fit_intercept"] == "true" else False
        lsvc_standardization = True if params["lsvc_standardization"] == "true" else False
        lsvc_aggregation_depth = int(params["lsvc_aggregation_depth"])
        cf = LinearSVC(
            maxIter = lsvc_max_iter,
            regParam = lsvc_reg_param,
            fitIntercept = lsvc_fit_intercept,
            standardization = lsvc_standardization,
            aggregationDepth = lsvc_aggregation_depth
        )
    elif type == 7: # Linear Regression
        l_rgs_max_iter = int(params["l_rgs_max_iter"])
        l_rgs_reg_param = float(params["l_rgs_reg_param"])
        cf = LinearRegression(
            maxIter = l_rgs_max_iter,
            regParam = l_rgs_reg_param,
        )
    elif type == 8: # Random Forest Regression
        rf_rgs_max_depth = int(params["rf_rgs_max_depth"])
        rf_rgs_num_trees = int(params["rf_rgs_num_trees"])
        cf = RandomForestRegressor(
            maxDepth = rf_rgs_max_depth,
            numTrees = rf_rgs_num_trees
        )
    elif type == 9: # Gradient-boosted Tree Regression
        gbt_rgs_max_depth = int(params["gbt_rgs_max_depth"])
        gbt_rgs_max_iter = int(params["gbt_rgs_max_iter"])
        cf = GBTRegressor(
            maxDepth = gbt_rgs_max_depth,
            maxIter = gbt_rgs_max_iter
        )

    return cf

# 获取当前项目下的数据集
@bp.route('get-dataset/', methods = ['POST'])
def getDataset():
    # 变量初始化 & 接收请求数据
    status = 0
    dataset_num = 0
    datasets = []
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # project_id
    project_id = received_data['id']

    # 获取数据集数量
    db = get_db()
    c = db.cursor()
    t = (project_id,)
    c.execute("SELECT COUNT(*) from dataset WHERE project_id=?", t)
    res = c.fetchone()
    dataset_num = tuple(res)[0]

    # 获取数据集
    if dataset_num == 0:
        datasets = []
    else:
        tmp = {
            "id": 0,
            "name": ""
        }
        for dataset in c.execute("SELECT id, name FROM dataset WHERE project_id=?", t):
            tmp["id"] = tuple(dataset)[0]
            tmp["name"] = tuple(dataset)[1]
            datasets.append(tmp.copy())

    status = 1

    # 接口约定：
    # status: 状态码
    # dataset_num: 数据集数量
    # datasets: 数据集数组
    #     [{"id": "...", "name": "..."}]
    return {
        "status": status,
        "dataset_num": dataset_num,
        "datasets": datasets
    }

# 获取数据集的列
@bp.route('get-columns/', methods = ['POST'])
def getColumns():
    # 变量初始化 & 接收请求数据
    status = 0
    columns = []
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # dataset_id
    dataset_id = received_data['dataset_id']

    # 获取 dataframe
    df = getDF(dataset_id)

    # 获取 columns
    tmp = {
        "index": 0,
        "name": "",
    }
    for i, column in enumerate(df.columns):
        tmp["index"] = i
        tmp["name"] = column
        columns.append(tmp.copy())

    status = 1

    # 接口约定：
    # status: 状态码
    # columns: 列
    return {
        "status": status,
        "columns": columns
    }

# 训练模型
@bp.route('train/', methods = ['POST'])
def train():
    # 变量初始化 & 接收请求数据
    status = 0
    results = {}
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # dataset_id, label, features, problem_type, classifier, classifier_params
    dataset_id = received_data['dataset_id']
    label = received_data['label']
    features = received_data['features']
    problem_type = received_data['problem_type']
    classifier = received_data['classifier']
    classifier_params = received_data['classifier_params']

    # 获取 dataframe
    df = getDF(dataset_id)

    # 数据预处理
    df = dataPreProcess(df, label, features, problem_type)

    # 划分训练集和测试集
    (trainingData, testData) = df.randomSplit([0.7, 0.3])

    # 训练
    cf = getClassifier(classifier, classifier_params)
    model = cf.fit(trainingData)

    # 把训练好的模型分别作用在训练集和测试集上
    predict_train = model.transform(trainingData)
    predict_test = model.transform(testData)

    # 模型评估
    if problem_type == 1: # 二分类问题
        evaluator = BinaryClassificationEvaluator( # 计算 area under ROC
            rawPredictionCol = "rawPrediction",
            labelCol = "label"
        )
        results["train_roc"] = evaluator.evaluate(predict_train)
        results["test_roc"] = evaluator.evaluate(predict_test)

        evaluator = BinaryClassificationEvaluator( # 计算 area under PR
            rawPredictionCol = "rawPrediction",
            labelCol = "label",
            metricName = "areaUnderPR"
        )
        results["train_pr"] = evaluator.evaluate(predict_train)
        results["test_pr"] = evaluator.evaluate(predict_test)

    elif problem_type == 2: # 多分类问题
        evaluator = MulticlassClassificationEvaluator( # 计算 acc
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'accuracy'
        )
        results["train_acc"] = evaluator.evaluate(predict_train)
        results["test_acc"] = evaluator.evaluate(predict_test)

        evaluator = MulticlassClassificationEvaluator( # 计算 f1 score
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'f1'
        )
        results["train_f1"] = evaluator.evaluate(predict_train)
        results["test_f1"] = evaluator.evaluate(predict_test)
    elif problem_type == 3: # 回归问题
        evaluator = RegressionEvaluator( # 计算 rmse
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'rmse'
        )
        results["train_rmse"] = evaluator.evaluate(predict_train)
        results["test_rmse"] = evaluator.evaluate(predict_test)

        evaluator = RegressionEvaluator( # 计算 mae
            labelCol = "label",
            predictionCol = "prediction",
            metricName = 'mae'
        )
        results["train_mae"] = evaluator.evaluate(predict_train)
        results["test_mae"] = evaluator.evaluate(predict_test)

    status = 1

    # 接口约定：
    # status: 状态码
    # results: 评估结果
    return {
        "status": status,
        "results": results
    }

# 保存模型
@bp.route('save/', methods = ['POST'])
def save():
    # 变量初始化 & 接收请求数据
    status = 0
    received_data = request.get_json(silent=True)

    # logging
    logging.info("received_data: " + json.dumps(received_data))

    # 应该接收到有:
    # dataset_id, label, features, problem_type, classifier
    # classifier_params, model_name, model_description
    dataset_id = received_data['dataset_id']
    label = received_data['label']
    features = received_data['features']
    problem_type = received_data['problem_type']
    classifier = received_data['classifier']
    classifier_params = received_data['classifier_params']
    model_name = received_data['model_name']
    model_description = received_data['model_description']

    # 获取 dataframe
    df = getDF(dataset_id)

    # 数据预处理
    df = dataPreProcess(df, label, features, problem_type)

    # 训练
    cf = getClassifier(classifier, classifier_params)
    model = cf.fit(df)

    db = get_db()
    c = db.cursor()

    #####################
    # 在 hdfs 中保存模型 #
    #####################

    # 通过 dataset_id 获取 project_id
    t = (dataset_id,)
    c.execute("SELECT project_id FROM dataset WHERE id=?", t)
    res = c.fetchone()
    project_id = tuple(res)[0]

    # 通过 project_id 获取 username
    t = (project_id,)
    c.execute("SELECT username FROM project WHERE id=?", t)
    res = c.fetchone()
    username = tuple(res)[0]

    # 拼接出 model_path 并保存模型
    model_path = "/user/" + username + "/" + str(project_id) + "/" + urllib.parse.quote(model_name)
    model.save(model_path)

    ####################
    # 在数据库中进行记录 #
    ####################

    t = (project_id, model_name, model_description, classifier, dataset_id, int(label))
    c.execute("INSERT INTO model (project_id, name, description, type, dataset_id, label_index) VALUES (?, ?, ?, ?, ?, ?)", t)
    db.commit()
    close_db()

    status = 1

    # 接口约定：
    # status: 状态码
    return {
        "status": status
    }


import argparse
import shlex

import greenlet as gl
import numpy as np
import tensorflow as tf

from pong.mechanics.pong import Pong, S
from pong.mechanics import policies
from pong.mechanics import constants as c
from pong.utils import common
from pong.utils.tf_machinery import NeuralNetwork

class NNPolicy(policies.PolicyRight):

    @classmethod
    def create_from_commandline(cls, side, argv):
        parser = argparse.ArgumentParser()
        parser.add_argument("save_path")
        parser.add_argument("--name", "-n", default="NeuralNetwork")
        args = parser.parse_args(shlex.split(argv))
        return cls(side, **vars(args))

    def __init__(self, side, layers=[50, 50, 50], save_path=None,
        name="Neural"):
        super().__init__(side, name)
        self.save_path = save_path

        self.g = tf.Graph()
        with self.g.as_default():
            self.nn = NeuralNetwork(Pong.STATE_DIM, Pong.NUM_ACTIONS, layers)
            if save_path is not None:
                self.nn.load(save_path)

    def get_action_right(self, state, *args):
        return gl.getcurrent().parent.switch((self.side, state))

    def evaluate_all(self, states):
        return self.nn.predict_argmax(states)

class TargetNNPolicy(policies.PolicyRight):

    @classmethod
    def create_from_commandline(cls, side, argv):
        parser = argparse.ArgumentParser()
        parser.add_argument("save_path")
        parser.add_argument("--name", "-n", default="NeuralNetwork")
        args = parser.parse_args(shlex.split(argv))
        return cls(side, **vars(args))

    def __init__(self, side, layers=[50, 50, 50], save_path=None, name="Neural"):
        super().__init__(side, name)
        self.save_path = save_path

        self.g = tf.Graph()
        with self.g.as_default():
            self.nn = NeuralNetwork(Pong.STATE_DIM, 1, layers)
            if save_path is not None:
                self.nn.load(save_path)

    def get_action_right(self, state, *args):
        return gl.getcurrent().parent.switch((self.side, state))

    def evaluate_all(self, states):
        y = states[:, S.R_Y]
        low = y - 0.5 * c.HPL
        high = y + 0.5 * c.HPL

        T = self.nn.predict_raw(states)[:, 0]
        A = np.zeros(states.shape[0], np.int32)
        A[T < low] = c.A_DOWN
        A[T > high] = c.A_UP
        return A

class FastEvaluate:
    def __init__(self, l_pol, r_pol, disc=False):
        self.l_pol = l_pol
        self.r_pol = r_pol
        self.disc = disc

    def run_episode(self, *args):
        if self.disc:
            sim = Pong(random_positions=True, f=self.r_pol.discretization)
        else:
            sim = Pong(random_positions=True)

        while not sim.done:
            state = sim.get_state()
            l_a = self.l_pol.get_action(state)
            r_a = self.r_pol.get_action(state)
            sim.step(l_a, r_a)

        return sim.win_string()

    def run_episodes(self, n):
        eval_states = np.zeros((n, Pong.STATE_DIM), np.float32)

        threads = []
        for i in range(n):
            t = gl.greenlet(self.run_episode)
            threads.append(t)

        A = [None] * n
        alive = np.ones(n, np.bool)

        while alive.any():
            flags = {"l": np.zeros(n, np.bool), "r": np.zeros(n, np.bool)}

            for i in range(n):
                if alive[i]:
                    data = threads[i].switch(A[i])
                    if not threads[i].dead:
                        side, state = data
                        eval_states[i] = state
                        flags[side][i] = True
                    else:
                        alive[i] = False
                        self.score[data] += 1

            A = np.zeros(n, np.int32)
            if flags["l"].any():
                A[flags["l"]] = self.l_pol.evaluate_all(eval_states[flags["l"]])
            if flags["r"].any():
                A[flags["r"]] = self.r_pol.evaluate_all(eval_states[flags["r"]])

    def estimate(self, n):
        self.score = {"l": 0, "r": 0, "draw": 0}
        self.run_episodes(n)
        return self.score

policies.POLICIES["nn"] = NNPolicy
policies.POLICIES["targetnn"] = TargetNNPolicy


from typing import Any, Optional, Sequence, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange
from monai.networks.blocks.convolutions import Convolution
from monai.networks.blocks.segresnet_block import ResBlock
from monai.networks.blocks.upsample import UpSample
from monai.networks.layers.factories import Act
from monai.networks.layers.utils import get_norm_layer
from monai.utils import InterpolateMode, UpsampleMode

from . import layers

conv1x1 = layers.ddpm_conv1x1
conv3x3 = layers.ddpm_conv3x3
NIN = layers.NIN
default_init = layers.default_init
AttentionBlock = layers.AttentionBlock

class GaussianFourierProjection(nn.Module):
    """Gaussian Fourier embeddings for noise levels."""

    def __init__(self, embedding_size=256, scale=1.0, learnable=False):
        super().__init__()
        self.W = nn.Parameter(
            torch.randn(embedding_size) * scale, requires_grad=learnable
        )

    def forward(self, x):
        x_proj = x[:, None] * self.W[None, :] * 2 * np.pi
        # print("x:", x.shape, "x_proj:", x_proj.shape)
        return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)

########## Code for 3D brain reconstruction models ##############

def get_conv_layer(
    spatial_dims: int,
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    bias: bool = True,
    dilation: int = 1,
    init_scale=1.0,
):
    conv = Convolution(
        spatial_dims,
        in_channels,
        out_channels,
        strides=stride,
        kernel_size=kernel_size,
        bias=bias,
        conv_only=True,
        dilation=dilation,
    )

    conv.conv.weight.data = default_init(init_scale)(conv.conv.weight.data.shape)
    if bias:
        nn.init.zeros_(conv.conv.bias)

    return conv

def get_upsample_layer(
    spatial_dims: int,
    in_channels: int,
    upsample_mode: Union[UpsampleMode, str] = "nontrainable",
    interp_mode=InterpolateMode.LINEAR,
    scale_factor: int = 2,
):
    return UpSample(
        spatial_dims=spatial_dims,
        in_channels=in_channels,
        out_channels=in_channels,
        scale_factor=scale_factor,
        mode=upsample_mode,
        interp_mode=interp_mode,
        align_corners=None if interp_mode == InterpolateMode.NEAREST else False,
    )

def make_dense_layer(in_sz, out_sz):
    dense = nn.Linear(in_sz, out_sz)
    dense.weight.data = default_init()(dense.weight.shape)
    nn.init.zeros_(dense.bias)
    return dense

class ResBlockpp(nn.Module):
    """
    Modified ResBlock form Monai implementation
    [LINK]
    ResBlock employs skip connection and two convolution blocks and is used
    in SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization
    <https://arxiv.org/pdf/1810.11654.pdf>`_.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        norm: Union[Tuple, str],
        act: str = "swish",
        kernel_size: int = 3,
        dilation: int = 1,
        init_scale=0.0,
    ) -> None:
        """
        Args:
            spatial_dims: number of spatial dimensions, could be 1, 2 or 3.
            in_channels: number of input channels.
            norm: feature normalization type and arguments.
            kernel_size: convolution kernel size, the value should be an odd number. Defaults to 3.
            dilation: dilation size of kernel for larger receptive fields
            init_scale: variance scale used to initialize kernel params via ddpm initializer
        """

        super().__init__()

        if kernel_size % 2 != 1:
            raise AssertionError("kernel_size should be an odd number.")

        self.norm1 = get_norm_layer(
            name=norm, spatial_dims=spatial_dims, channels=in_channels
        )
        self.norm2 = get_norm_layer(
            name=norm, spatial_dims=spatial_dims, channels=in_channels
        )
        # self.act = Act["mish"](inplace=True)
        self.act = Act[act]()

        # Following convention of the BigGAN blocks above
        # first conv uses default init scale, second uses config
        self.conv1 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            dilation=dilation,
            kernel_size=kernel_size,
        )
        self.conv2 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            dilation=dilation,
            init_scale=init_scale,
            kernel_size=kernel_size,
        )

    def forward(self, x):
        identity = x

        x = self.norm1(x)
        x = self.act(x)
        x = self.conv1(x)

        x = self.norm2(x)
        x = self.act(x)
        x = self.conv2(x)

        x += identity
        # x /= torch.sqrt(torch.tensor(2.0, requires_grad=False))

        return x

class SegResBlockpp(nn.Module):
    """
    ResBlock employs skip connection and two convolution blocks and is used
    in SegResNet based on `3D MRI brain tumor segmentation using autoencoder regularization
    <https://arxiv.org/pdf/1810.11654.pdf>`_.
    """

    def __init__(
        self,
        in_channels: int,
        norm: Union[Tuple, str],
        spatial_dims: int = 3,
        act: str = "swish",
        kernel_size: int = 3,
        temb_dim: int = None,
        pre_conv: Any = None,
        attention_heads: int = 0,
        dilation: int = 1,
        resblock_pp: bool = False,
        jit: bool = False,
    ) -> None:
        super().__init__()

        self.pre_conv = pre_conv
        self.n_channels = in_channels

        if not resblock_pp:
            self.resblock = ResBlock(
                spatial_dims, in_channels, norm, kernel_size=kernel_size
            )
        else:
            self.resblock = ResBlockpp(
                spatial_dims,
                in_channels,
                norm,
                dilation=dilation,
                kernel_size=kernel_size,
                act=act,
            )

        if jit:
            # print("Jitting resblock")
            self.resblock = torch.jit.script(self.resblock)
        self.attention = attention_heads

        if temb_dim is not None:
            self.dense = make_dense_layer(temb_dim, in_channels * 2)

        self.act = nn.SiLU()

        if attention_heads > 0:
            self.attn = AttentionBlock(channels=in_channels, num_heads=attention_heads)

    def forward(self, x, temb=None):
        if self.pre_conv is not None:
            x = self.pre_conv(x)

        x = self.resblock(x)

        # If time embedding provided
        # Conditioning is acheived via time embedding
        if temb is not None:
            cond_info = self.dense(self.act(temb))[:, :, None, None, None]
            gamma, beta = torch.split(
                cond_info, (self.n_channels, self.n_channels), dim=1
            )
            x = x * gamma + beta

        if self.attention:
            x = self.attn(x)

        return x

class ResnetBlockBigGANpp(nn.Module):
    """
    BigGAN block adapted from song. Notably, the conditioning is done b/w convs
    norm - act - conv --> time_cond --> norm - act - drop - conv --> skip
    """

    def __init__(
        self,
        in_channels,
        kernel_size: int = 3,
        spatial_dims: int = 3,
        act: str = "swish",
        temb_dim: int = None,
        dropout: float = 0.0,
        init_scale: float = 0.0,
        pre_conv: Any = None,
    ):
        super().__init__()

        self.n_channels = in_channels
        self.pre_conv = pre_conv

        # print("IN_CHANNELS:", in_channels)
        self.norm_0 = get_norm_layer(
            name=("GROUP", {"num_groups": min(in_channels // 4, 32), "eps": 1e-6}),
            spatial_dims=spatial_dims,
            channels=in_channels,
        )

        self.conv_0 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=kernel_size,
        )

        if temb_dim is not None:
            self.dense = make_dense_layer(temb_dim, in_channels * 2)

        self.norm_1 = get_norm_layer(
            name=("GROUP", {"num_groups": min(in_channels // 4, 32), "eps": 1e-6}),
            spatial_dims=spatial_dims,
            channels=in_channels,
        )
        self.conv_1 = get_conv_layer(
            spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            init_scale=init_scale,
            kernel_size=kernel_size,
        )
        self.dropout = nn.Dropout(dropout)

        self.act = Act[act]()

    def forward(self, x, temb):
        if self.pre_conv is not None:
            # This is where the downsample happens
            x = self.pre_conv(x)

        h = self.act(self.norm_0(x))

        h = self.conv_0(h)
        # FiLM-like conditioning for each feature map via time embedding
        cond_info = self.dense(self.act(temb))[:, :, None, None, None]
        gamma, beta = torch.split(cond_info, (self.n_channels, self.n_channels), dim=1)
        h = h * gamma + beta

        h = self.act(self.norm_1(h))
        h = self.dropout(h)
        h = self.conv_1(h)

        x = x + h

        return x

class AttentionBlock3d(nn.Module):
    """Channel-wise 3D self-attention block."""

    def __init__(
        self, channels, num_head_channels=None, skip_scale=False, init_scale=0.1
    ):
        super().__init__()
        torch.random.manual_seed(42)
        self.norm = nn.GroupNorm(num_groups=8, num_channels=channels, eps=1e-6)
        self.qkv = nn.Conv3d(channels, channels * 3, kernel_size=1)
        self.proj = nn.Conv3d(channels, channels, kernel_size=1)
        self.spatial_flatten = Rearrange(pattern="b c h w d -> b c (h w d)")
        self.scale = int(channels) ** (-0.5)
        self.skip_scale = np.sqrt(2.0) ** -1 if skip_scale else 1.0

        # self.num_heads = channels // num_head_channels if num_head_channels is not None else 1
        # self.scale = 1 / np.sqrt(channels / self.num_heads)

        # Initialize weights
        self.qkv.weight.data = default_init(init_scale)(self.qkv.weight.data.shape)
        self.proj.weight.data = default_init(init_scale)(self.proj.weight.data.shape)
        nn.init.zeros_(self.qkv.bias)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x):
        B, C, H, W, D = x.shape
        q, k, v = self.qkv(self.norm(x)).chunk(3, dim=1)

        q = self.spatial_flatten(q)
        k = self.spatial_flatten(k)
        v = self.spatial_flatten(v)

        with torch.cuda.amp.autocast(enabled=True, dtype=torch.float32):
            w = torch.einsum("b c q, b c k -> b q k", q, k) * self.scale
            w = F.softmax(w, dim=-1)
            # print("INSIDE ATTENTION BLOCK:", w.dtype)
        h = torch.einsum("b q k , b c k -> b c q", w, v)
        # print("OUTSIDE ATTENTION BLOCK:", h.dtype)
        h = torch.reshape(h, (B, C, H, W, D))
        h = self.proj(h)
        x = x + h
        x = x * self.skip_scale

        return x


import re
import random
import json

def get_domain(name):
    if name == 'object_division':
        return ObjectDivisionDomain()
    raise()

class ObjectDivisionDomain(object):
    def __init__(self):
        self.item_pattern = re.compile('^item([0-9])=([0-9\-])+$')

    def input_length(self):
        return 3

    def selection_length(self):
        return 6

    def generate_choices(self, inpt):
        cnts, _ = self.parse_context(inpt)

        def gen(cnts, idx=0, choice=[]):
            if idx >= len(cnts):
                left_choice = ['item%d=%d' % (i, c) for i, c in enumerate(choice)]
                right_choice = ['item%d=%d' % (i, n - c) for i, (n, c) in enumerate(zip(cnts, choice))]
                return [left_choice + right_choice]
            choices = []
            for c in range(cnts[idx] + 1):
                choice.append(c)
                choices += gen(cnts, idx + 1, choice)
                choice.pop()
            return choices
        choices = gen(cnts)
        choices.append(['<no_agreement>'] * self.selection_length())
        choices.append(['<disconnect>'] * self.selection_length())
        return choices

    def parse_context(self, ctx):
        cnts = [int(n) for n in ctx[0::2]]
        vals = [int(v) for v in ctx[1::2]]
        return cnts, vals

    def _to_int(self, x):
        try:
            return int(x)
        except:
            return 0

    def score_choices(self, choices, ctxs):
        assert len(choices) == len(ctxs)
        # print('choices = {}'.format(choices))
        # print('ctxs = {}'.format(ctxs))
        cnts = [int(x) for x in ctxs[0][0::2]]
        agree, scores = True, [0 for _ in range(len(ctxs))]
        for i, n in enumerate(cnts):
            for agent_id, (choice, ctx) in enumerate(zip(choices, ctxs)):
                # taken = self._to_int(choice[i+3][-1])
                taken = self._to_int(choice[i][-1])
                n -= taken
                scores[agent_id] += int(ctx[2 * i + 1]) * taken
            agree = agree and (n == 0)
        return agree, scores

class ContextGenerator(object):
    """Dialogue context generator. Generates contexes from the file."""
    def __init__(self, context_file):
        self.ctxs = []
        with open(context_file, 'r') as f:
            ctx_pair = []
            for line in f:
                ctx = line.strip().split()
                ctx_pair.append(ctx)
                if len(ctx_pair) == 2:
                    self.ctxs.append(ctx_pair)
                    ctx_pair = []

    def sample(self):
        return random.choice(self.ctxs)

    def iter(self, nepoch=1):
        for e in range(nepoch):
            random.shuffle(self.ctxs)
            for ctx in self.ctxs:
                yield ctx

    def total_size(self, nepoch):
        return nepoch*len(self.ctxs)

class ContextGeneratorEval(object):
    """Dialogue context generator. Generates contexes from the file."""
    def __init__(self, context_file):
        self.ctxs = []
        with open(context_file, 'r') as f:
            ctx_pair = []
            for line in f:
                ctx = line.strip().split()
                ctx_pair.append(ctx)
                if len(ctx_pair) == 2:
                    self.ctxs.append(ctx_pair)
                    ctx_pair = []

class TaskGoalGenerator(object):
    def __init__(self, goal_file):
        self.goals = []
        data = json.load(open(goal_file))
        for key, raw_dlg in data.items():
            self.goals.append((key, raw_dlg['goal']))

    def sample(self):
        return random.choice(self.goals)

    def iter(self, nepoch=1):
        for e in range(nepoch):
            random.shuffle(self.goals)
            for goal in self.goals:
                yield goal


from typing import List, Optional

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def constructMaximumBinaryTree(self, nums: List[int]) -> Optional[TreeNode]:

        if len(nums) == 0:
            return None

        tree = TreeNode()

        max_index: int = self.find_max(nums)
        tree.val = nums[max_index]
        tree.left = self.constructMaximumBinaryTree(nums[:max_index])
        if max_index + 1 <= len(nums) - 1:
            tree.right = self.constructMaximumBinaryTree(nums[max_index + 1:])
        else:
            tree.right = None
        return tree

    def find_max(self, nums: List[int]) -> int:
        max_value = float('-inf')
        max_index = 0
        for index, value in enumerate(nums):
            if value > max_value:
                max_value = value
                max_index = index
        return max_index


from dataclasses import dataclass, replace, field
from pprint import pprint
import abc
import typing
import torch

@dataclass(frozen=True)
class SequenceVariableValue:
    type = "sequence"
    value: torch.Tensor

@dataclass(frozen=True)
class MatrixVariableValue:
    type = "matrix"
    value: torch.Tensor

Variable = SequenceVariableValue | MatrixVariableValue

@dataclass(frozen=True)
class ProgramState:
    tokens: torch.Tensor
    sequence_variables: dict[str, SequenceVariableValue] = field(default_factory=dict)
    matrix_variables: dict[str, MatrixVariableValue] = field(default_factory=dict)

    def assign_sequence(
        self, variable_name: str, value: SequenceVariableValue
    ) -> "ProgramState":
        assert (
            variable_name not in self.sequence_variables
            and variable_name not in self.matrix_variables
        )
        return replace(
            self, sequence_variables={**self.sequence_variables, variable_name: value}
        )

    def assign_matrix(
        self, variable_name: str, value: MatrixVariableValue
    ) -> "ProgramState":
        assert (
            variable_name not in self.sequence_variables
            and variable_name not in self.matrix_variables
        )
        return replace(
            self, matrix_variables={**self.matrix_variables, variable_name: value}
        )

T = typing.TypeVar("T", bound=(SequenceVariableValue | MatrixVariableValue))

class Expression(abc.ABC, typing.Generic[T]):
    @abc.abstractmethod
    def evaluate(self, state: ProgramState) -> T:
        pass

SequenceExpression = Expression[SequenceVariableValue]
MatrixExpression = Expression[MatrixVariableValue]

Predicate = typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor]

def eq_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a == b

def lt_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a < b

def gt_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a > b

def leq_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a <= b

def geq_predicate(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a >= b

@dataclass(frozen=True)
class SelectExpression(MatrixExpression):
    k: SequenceExpression
    q: SequenceExpression
    p: Predicate

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=self.p(
                self.k.evaluate(state).value.unsqueeze(0),
                self.q.evaluate(state).value.unsqueeze(1),
            )
        )

@dataclass(frozen=True)
class SelectorWidthExpression(SequenceExpression):
    m: MatrixExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(self.m.evaluate(state).value.sum(dim=1))

@dataclass(frozen=True)
class NotSequenceExpression(SequenceExpression):
    s: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(~self.s.evaluate(state).value)

@dataclass(frozen=True)
class IndicatorExpression(SequenceExpression):
    s: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue((self.s.evaluate(state).value > 0).float())

@dataclass
class NotMatrixExpression(MatrixExpression):
    m: MatrixExpression

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(value=~self.m.evaluate(state).value)

@dataclass
class AndMatrixExpression(MatrixExpression):
    a: MatrixExpression
    b: MatrixExpression

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=self.a.evaluate(state).value & self.b.evaluate(state).value
        )

@dataclass(frozen=True)
class OrSequenceExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.a.evaluate(state).value | self.b.evaluate(state).value
        )

@dataclass(frozen=True)
class AndSequenceExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.a.evaluate(state).value & self.b.evaluate(state).value
        )

class TrueSequenceExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=torch.ones(state.tokens.shape[0], dtype=torch.bool)
        )

class TrueMatrixExpression(MatrixExpression):
    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=torch.ones(
                state.tokens.shape[0], state.tokens.shape[0], dtype=torch.bool
            )
        )

class TokensExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(value=state.tokens)

class IndicesExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(value=torch.arange(state.tokens.shape[0]))

@dataclass
class SequenceReferenceExpression(SequenceExpression):
    variable_name: str

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return state.sequence_variables[self.variable_name]

@dataclass
class MatrixReferenceExpression(MatrixExpression):
    variable_name: str

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return state.matrix_variables[self.variable_name]

Token = float

@dataclass(frozen=True)
class LiteralSequenceExpression(SequenceExpression):
    literal_value: Token

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=torch.full((state.tokens.shape[0],), self.literal_value)
        )

class LastColumnExpression(SequenceExpression):
    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=torch.full((state.tokens.shape[0],), state.tokens.shape[0] - 1)
        )

@dataclass(frozen=True)
class LiteralMatrixExpression(MatrixExpression):
    literal_value: bool

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=torch.full(
                (state.tokens.shape[0], state.tokens.shape[0]), self.literal_value
            )
        )

@dataclass(frozen=True)
class AggregateExpression(SequenceExpression):
    m: MatrixExpression
    s: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        s_value = self.s.evaluate(state).value
        m_value = self.m.evaluate(state).value
        mult = s_value * m_value
        agg = mult.sum(dim=1)
        avg = agg / m_value.sum(dim=1)
        return SequenceVariableValue(value=avg)

@dataclass(frozen=True)
class CompareSequencesExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression
    p: Predicate

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.p(
                self.a.evaluate(state).value,
                self.b.evaluate(state).value,
            )
        )

@dataclass(frozen=True)
class CompareMatricesExpression(MatrixExpression):
    a: MatrixExpression
    b: MatrixExpression
    p: Predicate

    def evaluate(self, state: ProgramState) -> MatrixVariableValue:
        return MatrixVariableValue(
            value=self.p(
                self.a.evaluate(state).value,
                self.b.evaluate(state).value,
            )
        )

@dataclass(frozen=True)
class SubtractSequencesExpression(SequenceExpression):
    a: SequenceExpression
    b: SequenceExpression

    def evaluate(self, state: ProgramState) -> SequenceVariableValue:
        return SequenceVariableValue(
            value=self.a.evaluate(state).value - self.b.evaluate(state).value
        )

@dataclass(frozen=True)
class Statement:
    assignment_name: str
    expression: SequenceExpression | MatrixExpression

@dataclass(frozen=True)
class Program:
    statements: tuple[Statement, ...]

    def invoke(self, tokens: torch.Tensor):
        state = ProgramState(tokens=tokens)
        for statement in self.statements:
            result = statement.expression.evaluate(state)
            if isinstance(result, SequenceVariableValue):
                state = state.assign_sequence(statement.assignment_name, result)
            elif isinstance(result, MatrixVariableValue):
                state = state.assign_matrix(statement.assignment_name, result)
            else:
                raise NotImplementedError()
        print(result)

double_histogram_program = Program(
    statements=(
        Statement(
            "same_tok",
            SelectExpression(TokensExpression(), TokensExpression(), eq_predicate),
        ),
        Statement(
            "hist", SelectorWidthExpression(MatrixReferenceExpression("same_tok"))
        ),
        Statement(
            "is_prev",
            SelectExpression(IndicesExpression(), IndicesExpression(), lt_predicate),
        ),
        Statement(
            "has_prev",
            CompareSequencesExpression(
                SelectorWidthExpression(
                    AndMatrixExpression(
                        MatrixReferenceExpression("is_prev"),
                        MatrixReferenceExpression("same_tok"),
                    ),
                ),
                LiteralSequenceExpression(0),
                gt_predicate,
            ),
        ),
        Statement(
            "first", NotSequenceExpression(SequenceReferenceExpression("has_prev"))
        ),
        Statement(
            "same_count",
            SelectExpression(
                SequenceReferenceExpression("hist"),
                SequenceReferenceExpression("hist"),
                eq_predicate,
            ),
        ),
        Statement(
            "first_mat",
            SelectExpression(
                SequenceReferenceExpression("first"),
                TrueSequenceExpression(),
                eq_predicate,
            ),
        ),
        Statement(
            "same_count_reprs",
            AndMatrixExpression(
                MatrixReferenceExpression("same_count"),
                MatrixReferenceExpression("first_mat"),
            ),
        ),
        Statement(
            "hist2",
            SelectorWidthExpression(MatrixReferenceExpression("same_count_reprs")),
        ),
    )
)

double_histogram_program.invoke(torch.tensor([1, 1, 1, 2, 2, 3, 3, 4, 5, 6]))

def frac_prevs(sop: SequenceExpression, val: Token) -> SequenceExpression:
    prevs = SelectExpression(IndicesExpression(), IndicesExpression(), leq_predicate)
    return AggregateExpression(
        prevs,
        IndicatorExpression(
            CompareSequencesExpression(
                sop, LiteralSequenceExpression(val), eq_predicate
            )
        ),
    )

def pair_balance(open_token: Token, close_token: Token) -> SequenceExpression:
    opens = frac_prevs(TokensExpression(), open_token)
    closes = frac_prevs(TokensExpression(), close_token)
    return SubtractSequencesExpression(opens, closes)

OPEN_A_TOKEN: Token = 1.0
CLOSE_A_TOKEN: Token = 2.0
OPEN_B_TOKEN: Token = 3.0
CLOSE_B_TOKEN: Token = 4.0

shuffle_dyck_2_program = Program(
    statements=(
        Statement("bal1", pair_balance(OPEN_A_TOKEN, CLOSE_A_TOKEN)),
        Statement("bal2", pair_balance(OPEN_B_TOKEN, CLOSE_B_TOKEN)),
        Statement(
            "negative",
            OrSequenceExpression(
                CompareSequencesExpression(
                    SequenceReferenceExpression("bal1"),
                    LiteralSequenceExpression(0),
                    lt_predicate,
                ),
                CompareSequencesExpression(
                    SequenceReferenceExpression("bal2"),
                    LiteralSequenceExpression(0),
                    lt_predicate,
                ),
            ),
        ),
        Statement(
            "had_neg",
            CompareSequencesExpression(
                AggregateExpression(
                    TrueMatrixExpression(),
                    IndicatorExpression(SequenceReferenceExpression("negative")),
                ),
                LiteralSequenceExpression(0),
                gt_predicate,
            ),
        ),
        Statement(
            "select_last",
            SelectExpression(IndicesExpression(), LastColumnExpression(), eq_predicate),
        ),
        Statement(
            "end_0",
            AggregateExpression(
                MatrixReferenceExpression("select_last"),
                AndSequenceExpression(
                    CompareSequencesExpression(
                        SequenceReferenceExpression("bal1"),
                        LiteralSequenceExpression(0),
                        eq_predicate,
                    ),
                    CompareSequencesExpression(
                        SequenceReferenceExpression("bal2"),
                        LiteralSequenceExpression(0),
                        eq_predicate,
                    ),
                ),
            ),
        ),
        Statement(
            "shuffle_dyck_2",
            AndSequenceExpression(
                CompareSequencesExpression(
                    SequenceReferenceExpression("end_0"),
                    LiteralSequenceExpression(1),
                    eq_predicate,
                ),
                NotSequenceExpression(SequenceReferenceExpression("had_neg")),
            ),
        ),
    )
)

shuffle_dyck_2_program.invoke(
    torch.tensor([OPEN_A_TOKEN, OPEN_B_TOKEN, CLOSE_A_TOKEN, CLOSE_B_TOKEN])
)

shuffle_dyck_2_program.invoke(
    torch.tensor(
        [OPEN_A_TOKEN, OPEN_B_TOKEN, CLOSE_A_TOKEN, CLOSE_B_TOKEN, CLOSE_B_TOKEN]
    )
)


from tkinter import *
import random
root=Tk()
root.title("Snake_water&gun")
root["bg"]="#eff"
root.geometry("1000x1000")

comp_list=["Snake","Water","Gun"]
def func(event):
	global var
	user=event.widget.cget("text")
	comp=random.choice(comp_list)
	txt=f"You selected {user}\nComputer selected {comp}"
	if user==comp:
		var.set(f"Tie!\n{txt}")
	elif (user=="Snake" and comp=="Water") or (user=="Water" and comp=="Gun") or (user=="Gun" and comp=="Snake"):
		var.set(f"You won!\n{txt}")
	else:
		var.set(f"You lost!\n{txt}")

	lbl.update()

frame1=Frame(root,bg="#eff")
frame1.pack(side=TOP,fill=X)

var=StringVar()
var.set("Let's Start")
lbl=Label(frame1, textvariable=var,bg="#eff")
lbl.pack(pady=(300,0))

frame2=Frame(root,bg="#eff")
frame2.pack(side=BOTTOM,ipady=180)

btn_list=["Snake","Water","Gun"]
for i in btn_list:
	btn=Button(frame2,text=i,bg="#ffb6c1", activebackground="#38aaf8")
	btn.pack(side=LEFT,padx=20)
	btn.bind("<Button-1>",func)

root.mainloop()


#!/usr/local/bin/python
# -*- coding: utf-8 -*-

import math
import numpy as np
import sys
sys.path.append("../")
from tensorflow2caffe import Tensorflow2Caffe

class Tensorflow2CaffeTransformerLstm(Tensorflow2Caffe):
    def __init__(self,
            tensorflow_model_path, caffe_model_path_prefix, caffe_model_name,
            max_seq_length, max_decode_length,
            encoder_params, decoder_params,
            use_small_word_list=False, max_candidate_size=0,
            check=False, calc=False):
        Tensorflow2Caffe.__init__(self, tensorflow_model_path,
            caffe_model_path_prefix, caffe_model_name, check, calc)
        self.scopes[0] = "seq2seq_model"
        self.encoder_params = encoder_params
        self.decoder_params = decoder_params
        self.max_seq_length = max_seq_length
        self.max_decode_length = max_decode_length
        self.use_small_word_list = use_small_word_list
        self.max_candidate_size = max_candidate_size
        if (self.use_small_word_list and self.max_candidate_size == 0):
            self.max_candidate_size = self.max_seq_length * 50 + 2000
        self.encoder_outputs = {}

    @staticmethod
    def default_encoder_params():
        return {
            "num_units": 512,
            "num_layers": 6,
            "layer.preprocess": "n",
            "layer.postprocess": "da",
            "ffn.num_units": 2048,
            "ffn.activation": "relu", # relu or swish
            "attention.num_heads": 8,
            "attention.branch": False,
            "attention.relpos": 0,
            "dropout_rate": 0.1,
            "position.enable": True,
            "position.combiner_fn": "tensorflow.add",
            "initializer":  "uniform_unit_scaling",
            "init_scale": 1.0,
            "share_level": 1 # every 2 layers share the same params
        }

    @staticmethod
    def default_decoder_params():
        return {
            "num_layers": 6,
            "num_units": 512,
            "layer.preprocess": "n",
            "layer.postprocess": "da",
            "attention.self_average": False,
            "attention.num_heads": 8,
            "attention.branch": False,
            "attention.relpos": 0,
            "ffn.num_units": 2048,
            "ffn.activation": "relu",
            "dropout_rate": 0.1,
            "position.enable": True,
            "position.combiner_fn": "tensorflow.add",
            "position.max_length": 1000,
            "decode_length_factor": 2.,
            "flex_decode_length": True,
            "initializer": "uniform_unit_scaling",
            "init_scale": 1.0,
            "attention.weighted_avg": False,
            "forget_bias": 1.0,
            "rnn.cell_type": "lstm",
            "sum_att": False
        }

    def ffn(self, x, output_name_prefix, scope_id, activation="relu"):
        self.scopes[scope_id] = "ffn"

        dense_name_1 = output_name_prefix + "_ffn_conv1"
        self.extract_dense(x, dense_name_1, scope_id+1, "conv1")

        activation_name = output_name_prefix + "_ffn_act"
        activation_support = False
        if (activation == "relu"):
            activation_support = True
            self.add_relu(dense_name_1, activation_name)
        if (activation == "swish"):
            activation_support = True
            self.add_swish(dense_name_1, activation_name)
        if (not activation_support):
            print("[ERROR] unsupported FFN activation %s" % (activation))
            exit(0)

        dense_name_2 = output_name_prefix + "_ffn_conv2"
        self.extract_dense(activation_name, dense_name_2, scope_id+1, "conv2")
        return dense_name_2

    def additive_attention(self, q,
                           k,
                           v,
                           mask,
                           attention_mask,
                           output_name_prefix,
                           name=None):
        print("[ERROR] unsupported additive attention")
        exit(0)

    def dot_product_attention(self, q,
                              k,
                              v,
                              mask,
                              attention_mask,
                              output_name_prefix,
                              edge_k=None,
                              edge_v=None,
                              name=None):
        if (edge_k is not None):
            sum_name = output_name_prefix + "_dot_ek"
            k = self.add_sum([k, edge_k], sum_name)
        # query * key
        query_key_name = output_name_prefix + "_dot_qk"
        self.add_matmul(q, k, query_key_name)

        if (mask is not None):
            scores = output_name_prefix + "_dot_scores"
            self.add_prod([query_key_name, mask], scores)
            query_key_name = output_name_prefix + "_dot_scores_mask"
            self.add_sum([scores, attention_mask], query_key_name)

        # softmax
        scores_normalized = output_name_prefix + "_dot_score_norm"
        self.add_softmax(query_key_name, scores_normalized, 3)

        if edge_v is not None:
            sum_name = output_name_prefix + "_dot_ev"
            v = self.add_sum([v, edge_v], sum_name)
        context = output_name_prefix + "_dot_cont"
        self.add_matmul(scores_normalized, v, context)

        return scores_normalized, context

    def multihead_attention(self, query,
                        memory,
                        mask,
                        attention_mask,
                        key_depth,
                        value_depth,
                        output_depth,
                        num_heads,
                        sequence_length,
                        output_name_prefix,
                        scope_id,
                        name=None,
                        cache=None,
                        branch=False,
                        filter_depth=None,
                        activation="relu",
                        relpos=0,
                        sum_att=False,
                        **kwargs):
        self.scopes[scope_id] = "multihead_attention"
        key_depth_per_head = key_depth // num_heads
        value_depth_per_head = value_depth // num_heads
        if memory is None:
            query_name = output_name_prefix + "_multihead_q"
            key_name = output_name_prefix + "_multihead_k"
            value_name = output_name_prefix + "_multihead_v"
            self.extract_denses(query, [query_name, key_name, value_name], [key_depth, key_depth, value_depth], scope_id+1, "qkv")
            key_reshape_name   = key_name + "_r"
            value_reshape_name = value_name + "_r"
            self.add_reshape(key_name,   key_reshape_name,   [self.batch, -1, num_heads, key_depth_per_head])
            self.add_reshape(value_name, value_reshape_name, [self.batch, -1, num_heads, value_depth_per_head])
            key_transpose_name   = key_name + "_t"
            value_transpose_name = value_name + "_t"
            self.add_transpose(key_reshape_name,   key_transpose_name,   [0, 2, 3, 1])
            self.add_transpose(value_reshape_name, value_transpose_name, [0, 2, 1, 3])
        else:
            query_name = output_name_prefix + "_multihead_q"
            self.extract_dense(query, query_name, scope_id+1, "q")
            #key_name = output_name_prefix + "_multihead_k"
            #value_name = output_name_prefix + "_multihead_v"
            #self.extract_denses(memory, [key_name, value_name], [key_depth, value_depth], scope_id+1, "kv")
            #key_name = memory["key"]
            #value_name = memory["value"]
            key_transpose_name = memory["key"]
            value_transpose_name = memory["value"]

        # reshape
        query_reshape_name = query_name + "_r"
        self.add_reshape(query_name, query_reshape_name, [self.batch, -1, num_heads, key_depth_per_head])

        # transpose
        query_transpose_name = query_name + "_t"
        self.add_transpose(query_reshape_name, query_transpose_name, [0, 2, 1, 3])

        edge_k = None
        edge_v = None
        if cache is not None:
            print("[ERROR] cache NOT_SUPPORTED")
            exit(0)

        query_scale_name = output_name_prefix + "_multihead_qs"
        self.add_power(query_transpose_name, query_scale_name, scale=1.0/math.sqrt(key_depth_per_head))

        if relpos > 0:
            print("[ERROR] relpos>0 NOT_SUPPORTED")
            exit(0)
        if sum_att:
            scores, x = self.additive_attention(
                query_scale_name, key_transpose_name, value_transpose_name, mask, attention_mask)
        else:
            scores, x = self.dot_product_attention(
                query_scale_name, key_transpose_name, value_transpose_name, mask, attention_mask, output_name_prefix,
                edge_k=edge_k, edge_v=edge_v)
        if branch:
            print("[ERROR] branch=True NOT_SUPPORTED")
            exit(0)
        else:
            # transpose
            x_t = output_name_prefix + "_multihead_out_t"
            self.add_transpose(x, x_t, [0, 2, 1, 3])
            # reshape
            x_r = output_name_prefix + "_multihead_out_r"
            #self.add_reshape(x_t, x_r, [self.batch, sequence_length, num_heads*value_depth_per_head])
            self.add_reshape(x_t, x_r, [self.batch, -1, num_heads*value_depth_per_head])
            # dense
            x = output_name_prefix + "_multihead_out_dense"
            self.extract_dense(x_r, x, scope_id+1, "output_transform")
        return scores, x

    def self_attention_sublayer(self, x, mask, attention_mask, num_units, num_heads, sequence_length, output_name_prefix, scope_id, memory=None, cache=None, branch=False,
                            filter_depth=None, activation="relu", relpos=0, sum_att=False):
        att_scores, x = self.multihead_attention(
            query=x,
            memory=memory,
            mask=mask,
            attention_mask=attention_mask,
            key_depth=num_units,
            value_depth=num_units,
            output_depth=num_units,
            num_heads=num_heads,
            sequence_length=sequence_length,
            output_name_prefix=output_name_prefix,
            scope_id=scope_id,
            cache=cache,
            branch=branch,
            filter_depth=filter_depth,
            activation=activation,
            relpos=relpos,
            sum_att=sum_att
        )
        return att_scores, x

    def layer_process(self, x, output_name_prefix, scope_id, y=None, mode=None):
        if not mode or mode == "none":
          return x

        index = 0
        for m in mode:
            if m == 'a':
                output_name = output_name_prefix + "_a" + str(index)
                x = self.add_sum([x, y], output_name)
            elif m == 'n':
                output_name = output_name_prefix + "_n" + str(index)
                x = self.extract_layer_norm(x, output_name, scope_id, ["layer_norm", "gamma", "beta"])
            elif m == 'd':
                print("[INFO] dropout")
            else:
                print("[ERROR] unknown layer process %s" % (m))
            index += 1
        return x

    def position_encoding(self, length, depth, output_name_prefix=None,
                          min_timescale=1,
                          max_timescale=1e4):
        positions = np.arange(length)
        depths = np.arange(depth)
        # correspond to log(10000^(1/(d-1)))
        log_timescale_increment = (
            math.log(max_timescale / min_timescale) / (depth - 1))
        # correspond to 1 / 10000^(i/(d-1)), i=0....d-1
        inv_timescales = min_timescale * np.exp(depths * -1 * log_timescale_increment)
        # pos / 10000^(i/(d-1))
        scaled_time = np.expand_dims(positions, 1) * np.expand_dims(inv_timescales, 0)
        # intead of using SIN and COS interleaved
        # it's  the same to first use SIN then COS
        # as they are applied to the same position
        position_embedding_weight = np.concatenate((np.sin(scaled_time), np.cos(scaled_time)), axis=1)

        output_name = "_position_dict"
        if (output_name_prefix is not None):
            output_name = output_name_prefix + output_name
        self.add_weight(output_name=output_name, weight=position_embedding_weight)
        return output_name

    def encode(self, inputs, sequence_mask, attention_mask, scope_id, output_name_prefix):
        num_units = self.encoder_params["num_units"]

        self.scopes[scope_id] = "target_space_emb"
        adder = output_name_prefix + "_emb"
        weight_name = output_name_prefix + "_target_space_emb"
        self.add_weight(weight_name, scope_id=scope_id+1)
        self.add_sum([inputs, weight_name], adder)

        if (self.encoder_params["position.enable"]):
            position_input_name = "nmt_positions"
            position_input_shape = [self.batch, self.max_seq_length]
            self.add_input(position_input_name, position_input_shape)
            #weight = np.array([[i for i in range(sequence_length)] * self.batch])
            #self.add_weight(position_input_name, weight=weight, data_type="INT32")

            position_embedding_dict_name = self.position_encoding(length=self.max_seq_length,
                                               depth=self.encoder_params["num_units"] // 2,
                                               output_name_prefix=output_name_prefix)
            position_embedding_name = output_name_prefix + "position_embedding"
            self.add_embedding(position_input_name,
                               position_embedding_dict_name,
                               position_embedding_name)

            if (self.encoder_params["position.combiner_fn"] != "tensorflow.add"):
                print("[ERROR] position embedding unsupported")
                exit(0)
            output_name = "we+pe"
            self.add_sum([adder, position_embedding_name], output_name)
            adder = output_name

        x = adder
        for i in range(self.encoder_params["num_layers"]):
            layer_idx = i // self.encoder_params["share_level"]
            self.scopes[scope_id] = "layer_" + str(layer_idx)
            self.scopes[scope_id+1] = "self_attention"
            output_name_prefix_new = output_name_prefix + "_layer" + str(i)

            # preprocess
            x_preprocess = self.layer_process(x,
                               output_name_prefix_new + "_pre1",
                               scope_id+2,
                               mode=self.encoder_params["layer.preprocess"])

            # attention
            _, y = self.self_attention_sublayer(x=x_preprocess,
                                  mask=sequence_mask,
                                  attention_mask=attention_mask,
                                  num_units=self.encoder_params["num_units"],
                                  num_heads=self.encoder_params["attention.num_heads"],
                                  sequence_length=self.max_seq_length,
                                  output_name_prefix=output_name_prefix_new,
                                  scope_id=scope_id+2,
                                  branch=self.encoder_params["attention.branch"],
                                  filter_depth=self.encoder_params["ffn.num_units"],
                                  activation=self.encoder_params["ffn.activation"],
                                  relpos=self.encoder_params["attention.relpos"],
                              )

            # post process
            x = self.layer_process(x,
                    output_name_prefix_new + "_post1",
                    scope_id+2,
                    y=y, mode=self.encoder_params["layer.postprocess"])

            # ffn
            self.scopes[scope_id+1] = "ffn"
            x_preprocess = self.layer_process(x,
                               output_name_prefix_new + "_pre2",
                               scope_id+2,
                               mode=self.encoder_params["layer.preprocess"])
            y = self.ffn(x_preprocess, output_name_prefix_new, scope_id+2, activation=self.encoder_params["ffn.activation"])
            x = self.layer_process(x,
                    output_name_prefix_new + "_post2",
                    scope_id+2,
                    y=y, mode=self.encoder_params["layer.postprocess"])

        outputs = self.layer_process(x,
                      output_name_prefix + "_att_post",
                      scope_id,
                      mode=self.encoder_params["layer.preprocess"])
        self.encoder_outputs["encoder_output"] = outputs
        return outputs

    def encoder_post_process(self, output_name_prefix):
        self.scopes[0] = "seq2seq_model"
        self.scopes[1] = "rnnformer_decoder"
        self.scopes[3] = "encdec_attention"
        self.scopes[4] = "multihead_attention"
        key_depth = self.decoder_params["num_units"]
        value_depth = self.decoder_params["num_units"]
        num_heads = self.decoder_params["attention.num_heads"]
        key_depth_per_head = key_depth // num_heads
        value_depth_per_head = value_depth // num_heads
        encoder_outputs = []
        for layer_idx in range(self.decoder_params["num_layers"]):
            self.scopes[2] = "layer_%d" % (layer_idx)
            key_name = output_name_prefix + "_layer" + str(layer_idx) + "_multihead_k"
            value_name = output_name_prefix + "_layer" + str(layer_idx) + "_multihead_v"
            memory = self.encoder_outputs["encoder_output"]
            self.extract_denses(memory, [key_name, value_name], [key_depth, value_depth], 5, ["kv", "kernel", "bias"])
            key_reshape_name   = key_name + "_r"
            value_reshape_name = value_name + "_r"
            self.add_reshape(key_name,   key_reshape_name,   [self.batch, -1, num_heads, key_depth_per_head])
            self.add_reshape(value_name, value_reshape_name, [self.batch, -1, num_heads, value_depth_per_head])
            key_transpose_name   = key_name + "_t"
            value_transpose_name = value_name + "_t"
            self.add_transpose(key_reshape_name,   key_transpose_name,   [0, 2, 3, 1])
            self.add_transpose(value_reshape_name, value_transpose_name, [0, 2, 1, 3])
            encoder_outputs.append({"key": key_transpose_name, "value": value_transpose_name})
        self.encoder_outputs["encoder_output"] = encoder_outputs

    def attention_ffn_block(self, inputs, encoder_mask, attention_mask, scope_id, output_name_prefix, state=None, position=None):
        x = inputs
        attentions = []
        state_cache = []
        for layer_idx in range(self.decoder_params["num_layers"]):
            self.scopes[scope_id] = "layer_%d" % (layer_idx)
            output_name_prefix_new = output_name_prefix + "_layer_" + str(layer_idx)

            # RNN sublayer
            self.scopes[scope_id+1] = "rnn_sublayer"
            cur_state = state[layer_idx]
            # Preprocess
            x_process = output_name_prefix_new + "_pre1"
            x = self.layer_process(x, x_process, scope_id+2, mode=self.decoder_params["layer.preprocess"])

            x = self.add_squeeze(x, axis=1, output_name=x+"_squeeze")

            y = output_name_prefix_new + "_cell"
            self.extract_rnn("LSTM", x, cur_state, y, scope_id+2, scope_name = "basic_lstm_cell")
            state_cache.append(cur_state)
            #Postprocess
            x_process = output_name_prefix_new + "_post1"
            x = self.layer_process(x, x_process, scope_id+2, y=y, mode=self.decoder_params["layer.postprocess"])
            x = self.add_expand_dims(x, axis=1, output_name=x+"_expand")

            # Encdec sublayer
            self.scopes[scope_id+1] = "encdec_attention"
            # Preprocess
            x_preprocess = output_name_prefix_new + "_pre2"
            x_preprocess = self.layer_process(x, x_preprocess, scope_id+2, mode=self.decoder_params["layer.preprocess"])
            # Cross attention
            att_scores, y = self.self_attention_sublayer(x=x_preprocess,
                                  mask=encoder_mask,
                                  attention_mask=attention_mask,
                                  num_units=self.decoder_params["num_units"],
                                  num_heads=self.decoder_params["attention.num_heads"],
                                  sequence_length=1,
                                  output_name_prefix=output_name_prefix_new,
                                  scope_id=scope_id+2,
                                  memory=self.encoder_outputs["encoder_output"][layer_idx],
                                  branch=self.decoder_params["attention.branch"],
                                  filter_depth=self.decoder_params["ffn.num_units"],
                                  activation=self.decoder_params["ffn.activation"],
                                  sum_att=self.decoder_params["sum_att"]
                              )
            # Post process
            x_process = output_name_prefix_new + "_post2"
            x = self.layer_process(x, x_process, scope_id+2, y=y, mode=self.decoder_params["layer.postprocess"])

            att_context = x

            if not self.decoder_params["attention.weighted_avg"]:
                print("[WARNING] unused attention scores")
                attentions.append(att_scores)
            else:
                print("[ERROR] unsupported attention weighted average")

            #if attention is None:
            #    attention = att_scores
            #else:
            #    if not self.decoder_params["attention.weighted_avg"]:
            #        output_name = attention + "add_socres"
            #        #attention = self.add_sum([attention, att_scores], output_name)
            #    else:
            #        print("[ERROR] unsupported attention weighted average")
            #        #attention = np.concatenate([attention, att_scores], axis = 1)

            # FFN sublayer
            self.scopes[scope_id+1] = "ffn"
            # Preprocess
            x_preprocess = output_name_prefix_new + "_pre3"
            x_preprocess = self.layer_process(x, x_preprocess, scope_id+2, mode=self.decoder_params["layer.preprocess"])
            # FFN
            y = self.ffn(x_preprocess, output_name_prefix_new, scope_id+2, activation=self.decoder_params["ffn.activation"])
            # Postprocess
            x_process = output_name_prefix_new + "_post3"
            x = self.layer_process(x, x_process, scope_id+2, y=y, mode=self.decoder_params["layer.postprocess"])
        state = state_cache

        #Preproecss
        x_process = output_name_prefix_new + "_pre4"
        outputs = self.layer_process(x, x_process, scope_id, mode=self.decoder_params["layer.preprocess"])

        attention = None
        if not self.decoder_params["attention.weighted_avg"]:
            attention = self.add_sum(attentions, "decoder_attention_sum")
            attention = self.add_reduce_mean(attention, axis=1, keep_dim=False, output_name="decoder_attention_mean")
            attention = self.add_power(attention, "decoder_attention_avg", scale=1.0/self.decoder_params["num_layers"])
        else:
            print("[ERROR] unsupported attention weighted average")

        return outputs, state, attention, att_context

    def add_projection(self, input_name, weight_name, output_name_prefix):
        matmul_name = output_name_prefix + "_matmul"
        self.add_matmul(input_name, weight_name, matmul_name, transpose_a=False, transpose_b=True)

        argmax_name = self.add_argmax(matmul_name, axis=-1, output_name=output_name_prefix+"_argmax")

        return argmax_name

    def extract_decoder(self, sequence_mask, attention_mask, max_decode_length, scope_id, output_name_prefix):
        # sos=1
        sos = "sos"
        weight = np.array([[1] * self.batch])
        self.add_weight(sos, weight=weight, data_type="INT32")

        negative_one = "negative_one"
        weight = np.array([[-1] * self.batch])
        self.add_weight(negative_one, weight=weight, data_type="INT32")

        decoder_start_name = output_name_prefix + "_words"
        decoder_start_shape = [self.batch, 1]
        self.add_memory(decoder_start_name, decoder_start_shape, data_type="INT32")
        self.add_copy(sos, 1, 1, 0,
                      decoder_start_name, 1, 1, 0,
                      1, output_name="init_decoder")

        position_input_name = output_name_prefix + "_position"
        position_input_shape = [self.batch, 1]
        self.add_memory(position_input_name, position_input_shape, data_type="INT32")
        self.add_copy(negative_one, 1, 1, 0,
                      position_input_name, 1, 1, 0,
                      1, output_name="init_decoder_position")

        zero = "zero"
        weight = np.array([[0]*self.batch])
        self.add_weight(zero, weight=weight, data_type="INT32")

        # eos=2
        eos = "eos"
        weight = np.array([[2]*self.batch])
        self.add_weight(eos, weight=weight, data_type="INT32")

        decoder_output_shape = [self.batch, max_decode_length]
        decoder_output = "decoder_output"
        self.add_memory(decoder_output, decoder_output_shape, data_type="INT32")

        decoder_attention_shape = [self.batch, max_decode_length, self.max_seq_length]
        decoder_attention = "decoder_attention"
        self.add_memory(decoder_attention, decoder_attention_shape, data_type="FLOAT32")

        position_embedding_dict_name = self.position_encoding(length=self.decoder_params["position.max_length"],
                                           depth=self.decoder_params["num_units"] // 2,
                                           output_name_prefix=output_name_prefix)
        state = []
        state_shape = [self.batch, self.encoder_params["num_units"]+self.decoder_params["num_units"]]
        for layer_idx in range(self.decoder_params["num_layers"]):
            state_name = output_name_prefix + "_layer" + str(layer_idx) + "_state"
            self.add_memory(state_name, state_shape, data_type="FLOAT32")
            state.append(state_name)
        sample_ids = decoder_start_name
        for step in range(max_decode_length):
            # whether to add caffe layer
            self.set_add_layer(step==0)

            position_input_name_new = position_input_name+"_add_one"
            self.add_power(position_input_name, position_input_name_new, scale=1, shift=1)
            self.add_copy(position_input_name_new, 1, 1, 0,
                          position_input_name, 1, 1, 0,
                          1, output_name="update_position")

            output_name_prefix_new = output_name_prefix + "_step" + str(step)
            cur_inputs = output_name_prefix_new + "_words_embedding"
            self.add_embedding(sample_ids, self.source_modality, cur_inputs)#, transpose=True)
            cur_inputs_scale = output_name_prefix_new + "_words_embedding" + "_s"
            self.add_power(cur_inputs, cur_inputs_scale, scale=(self.decoder_params['num_units'] ** 0.5))

            position_embedding_name = output_name_prefix_new + "_position_embedding"
            self.add_embedding(position_input_name, position_embedding_dict_name, position_embedding_name)

            output_name = output_name_prefix_new + "_embedding"
            cur_inputs_pos = self.add_sum([cur_inputs_scale, position_embedding_name], output_name)

            cell_outputs, state, attention, att_context = self.attention_ffn_block(inputs=cur_inputs_pos,
                                encoder_mask=sequence_mask,
                                attention_mask=attention_mask,
                                scope_id=scope_id,
                                output_name_prefix = output_name_prefix_new,
                                state=state,
                                position=step)
            self.add_copy(attention,
                          -1, -1, 0,
                          decoder_attention,
                          max_decode_length*self.max_seq_length, self.max_seq_length, 0,
                          -1,
                          output_name="copy_attention_to_global_buffer",
                          src_index_name=zero,
                          dst_index_name=position_input_name)

            current_ids = self.add_projection(cell_outputs, self.source_modality, output_name)
            self.add_copy(current_ids, 1, 1, 0,
                          sample_ids, 1, 1, 0,
                          1, output_name="copy_to_next_input")
            self.add_copy(current_ids,
                          1, 1, 0,
                          decoder_output,
                          max_decode_length, 1, 0,
                          1,
                          output_name="copy_word_to_global_buffer",
                          src_index_name=zero,
                          dst_index_name=position_input_name)
            status = output_name + "_check"
            self.add_check(current_ids, eos, "equal", status)
            self.add_repeat(max_decode_length-1, position_input_name_new, output_name="repeat", status_name=status)
            if (self.get_tensor(status)[0]):
                break;

        return self.get_tensor(decoder_output), self.get_tensor(decoder_attention)

    def generate(self, input=None):
        encoder_word_input_name = "nmt_words"
        encoder_word_input_shape = [self.batch, self.max_seq_length]
        self.add_input(encoder_word_input_name, encoder_word_input_shape)
        if (self.use_small_word_list):
            decoder_candidate_input_name = "nmt_candidates"
            decoder_candidate_input_shape = [self.batch, self.max_candidate_size]
            self.add_input(decoder_candidate_input_name, decoder_candidate_input_shape)
        self.set_input(input)

        self.scopes[1] = "source_modality"
        self.scopes[2] = "embedding"
        self.source_modality = "source_modality"
        self.add_weight(output_name=self.source_modality, scope_id=3)#, transpose=[1,0])

        encoder_embedding = "encoder_embedding"
        self.add_embedding(encoder_word_input_name, self.source_modality, encoder_embedding)#, transpose=True)
        encoder_embedding = self.add_power(encoder_embedding, encoder_embedding+"_s", scale=math.sqrt(self.encoder_params["num_units"]))

        if (self.use_small_word_list):
            small_words_embedding = "small_words_embedding"
            self.add_embedding(decoder_candidate_input_name, self.source_modality, small_words_embedding)#, transpose=True)
            self.source_modality = small_words_embedding
        mask_input_name = None
        encoder_attention_mask = None
        decoder_attention_mask = None
        #mask_input_name = "nmt_mask"
        #mask_input_shape = [self.batch, self.max_seq_length]
        #self.add_input(mask_input_name, mask_input_shape)

        self.scopes[1] = "transformer_encoder"
        #encoder_attention_mask = "encoder_attention_mask"
        #self.add_attention(mask_input_name, self.encoder_params['attention.num_heads'], self.max_seq_length, self.max_seq_length, encoder_attention_mask)
        encoders = self.encode(encoder_embedding, mask_input_name, encoder_attention_mask, 2, output_name_prefix="transformer")
        self.encoder_post_process("transformer_decoder")

        self.scopes[1] = "rnnformer_decoder"
        #decoder_attention_mask = "decoder_attention_mask"
        #self.add_attention(mask_input_name, self.decoder_params['attention.num_heads'], 1, self.max_seq_length, decoder_attention_mask)
        decoders = self.extract_decoder(mask_input_name, decoder_attention_mask, self.max_decode_length, 2, output_name_prefix="rnnformer")

        self.save_caffe_model()