#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import OrderedDict
from .tc import TC_Whitelist
from .utility import Utility
from .base import OperatorLayerBase


def hasTileSize(name):
    if ("sgemm" in name) or ("884gemm" in name) or ("hgemm" in name):
        return True
    else:
        return False


def ctaTile(name):
    name = name.split("_")
    name = list(filter(lambda x: "x" in x, name))
    name = list(filter(lambda x: "slice" not in x, name))
    assert (len(name) == 1)
    name = name[0].split("x")
    assert (len(name) == 2)
    name = list(map(int, name))
    return name[0], name[1]


class RNNCell(OperatorLayerBase):
    """
	This class supports RNNCell, LSTMCell and GRUCell.
	"""

    def __init__(self, d):
        marker = eval(d.argMarker[0])
        mod = marker['mod']
        op = marker['op']
        args = marker['args']

        self.marker = marker
        self.mod_ = mod
        self.op_ = op
        self.args = args

        self.name = d.name
        self.dir = d.dir
        self.sub = d.sub
        self.grid = d.grid

        assert (op == "forward")
        assert (mod in ["LSTMCell", "GRUCell", "RNNCell"])
        assert (len(args) in [2, 3])

        x, h = args[0], args[1]
        b1, ii = x['shape']
        b2, hh = h['shape']
        assert b1 == b2
        assert x['dtype'] == h['dtype']
        t = x['dtype']

        self.cell = mod
        self.inp = ii
        self.hid = hh
        self.b = b1
        self.type = t

        self.multiple = 1
        if self.cell == "LSTMCell":
            self.multiple = 4
        elif self.cell == "GRUCell":
            self.multiple = 3

        self.gemm = None
        self.m = None
        self.n = None
        self.k = None
        self.elems = 0

        self.bar()

    def params(self):
        if self.gemm is None:
            p = OrderedDict([('cell', self.cell), ('X', self.inp), ('H', self.hid), ('B', self.b), ('type', self.type)])
        else:
            assert self.m is not None
            assert self.n is not None
            assert self.k is not None
            p = OrderedDict([('gemm', self.gemm), ('M', self.m), ('N', self.n), ('K', self.k), ('type', self.type)])
        return p

    def tc(self):
        if "gemm" in self.name:
            if self.name in TC_Whitelist():
                return 1
            return 0
        else:
            return "-"

    def op(self):
        return self.op_

    def mod(self):
        return self.mod_

    def bytes(self):
        if self.gemm is not None:
            m, n, k, t = self.m, self.n, self.k, self.type
            b = (m * k + k * n + m * n) * Utility.typeToBytes(t)
        elif self.elems != 0:
            b = self.elems * Utility.typeToBytes(self.type)
        else:
            b = 0
        return b

    def flops(self):
        if self.gemm is not None:
            m, n, k = self.m, self.n, self.k
            f = 2 * m * n * k
        elif self.elems != 0:
            f = 0  #TODO
        else:
            f = 0
        return f

    def bar(self):
        cell = self.cell
        X = self.inp
        H = self.hid
        B = self.b
        t = self.type
        subseqId = self.sub
        direc = self.dir
        name = self.name
        grid = self.grid
        multiple = self.multiple

        if direc == "fprop":
            subseqId = subseqId % 3
            if subseqId == 0:  #layer gemm
                self.gemm = "layer"
                self.m = multiple * H
                self.n = B
                self.k = X
            elif subseqId == 1:  #recurrent gemm
                self.gemm = "recur"
                self.m = multiple * H
                self.n = B
                self.k = H
            else:
                layerGemmElems = multiple * H * B
                recurGemmElems = multiple * H * B
                cElems = H * B
                hElems = H * B
                totElems = layerGemmElems + recurGemmElems + 2 * cElems + hElems
                self.elems = totElems

        else:
            if ("gemm" in name) and hasTileSize(name):  #gemm
                #Get cta tile size
                tileX, tileY = ctaTile(name)
                #Get grid dimensions
                grid = grid.split(",")
                gridX, gridY, gridZ = map(lambda x: int(x), grid)

                gemmM = tileX * gridX
                gemmN = tileY * gridY

                if name[-3:] == "_nn":  # dgrad
                    if (gemmM == H):  # recurrent dgrad
                        #Ideally gemmN = B, but we have a limited set of tile sizes.
                        gemmN = B
                        gemmK = multiple * H

                        self.gemm = "recur"
                        self.m = gemmM
                        self.n = gemmN
                        self.k = gemmK

                    elif (gemmM == X):  # layer dgrad
                        #assert(gemmN % B == 0)
                        gemmK = multiple * H

                        self.gemm = "layer"
                        self.m = gemmM
                        self.n = gemmN
                        self.k = gemmK

                    else:
                        pass

                elif name[-3:] == "_nt":  #wgrad
                    if (gemmM == H):  #recurrent wgrad
                        assert (gemmN == multiple * H)
                        gemmK = B

                        self.gemm = "recur"
                        self.m = gemmM
                        self.n = gemmN
                        self.k = gemmK

                    elif (gemmM == X):  #layer wgrad
                        assert (gemmN == multiple * H)
                        gemmK = B

                        self.gemm = "layer"
                        self.m = gemmM
                        self.n = gemmN
                        self.k = gemmK

                    else:
                        pass
                else:
                    pass
            else:
                pass

        return
