#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import numpy as np

from coremltools import _logger as logger
from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
from coremltools.converters.mil.mil.passes.pass_registry import register_pass


@register_pass(namespace="tensorflow")
class expand_tf_lstm(AbstractGraphPass):
    """
    Expand tf_lstm_block_cell to fine-grained SSA ops following:

    xh = [x, h_prev]
    [i, ci, f, o] = xh * w + b
    f = f + forget_bias
    if not use_peephole:
      wci = wcf = wco = 0
    i = sigmoid(cs_prev .* wci + i)
    f = sigmoid(cs_prev .* wcf + f)
    ci = tanh(ci)
    cs = ci .* i + cs_prev .* f
    cs = clip(cs, cell_clip)
    o = sigmoid(cs * wco + o)
    co = tanh(cs)
    h = co .* o

    Inputs:

        prog: Program
    """
    def apply(self, prog):
        for f in prog.functions.values():
            _expand_tf_lstm_helper(f)


def _expand_tf_lstm_helper(block):
    # shallow copy hides changes on f.operations during the loop
    for op in list(block.operations):
        for b in op.blocks:
            _expand_tf_lstm_helper(b)

        if op.op_type == "tf_lstm_block_cell":
            _expand_tf_lstm_block_cell(op)
            logger.info("Expanding {} (op_type: {})".format(op.name, op.op_type))

        if op.op_type == "tf_lstm_block":
            # only cs, h are supported for now. Can be easily extended to other outputs at performance hit.
            i, cs, f, o, ci, co, h = op.outputs
            if all(
                map(lambda x: len(x.child_ops) == 0 and len(x.consuming_blocks) == 0,
                    (i, f, o, ci, co)
                )
            ):
                _expand_tf_lstm_block(op)
                logger.info("Expanding {} (op_type: {})".format(op.name, op.op_type))


def _lstm_cell_builder(op, x, h_prev, cs_prev):
    b = op.bias  # [4*hidden_dim]
    forget_bias = op.forget_bias.val  # python:float

    # xh = [x, h_prev]
    # xh shape: [b, input_dim+hidden_dim]
    xh = mb.concat(values=[x, h_prev], axis=-1)

    # w: [4*hidden_dim, input_dim + hidden_dim] (icfo layout)
    w = np.transpose(op.weight.val)
    # [i, ci, f, o] = xh * w + b. Shape is [b, 4*hidden_dim]
    icfo = mb.linear(x=xh, weight=w, bias=b)

    # i, ci, f, o shape: [b, hidden_dim]
    i, ci, f, o = mb.split(x=icfo, num_splits=4, axis=-1)
    if op.forget_bias.val != 0:
        f = mb.add(x=f, y=forget_bias)

    # note that .* means Hadamard product
    # i = sigmoid(cs_prev .* wci + i)
    # f = sigmoid(cs_prev .* wcf + f)
    if op.use_peephole.val:
        wci = op.weight_peep_i.val  # [hidden_dim]
        wcf = op.weight_peep_f.val  # [hidden_dim]

        x = mb.mul(x=cs_prev, y=wci)
        pre_i = mb.add(x=x, y=i)

        x = mb.mul(x=cs_prev, y=wcf)
        pre_f = mb.add(x=x, y=f)
    else:
        pre_i = i
        pre_f = f

    i = mb.sigmoid(x=pre_i)
    f = mb.sigmoid(x=pre_f)
    ci = mb.tanh(x=ci)

    # cs = ci .* i + cs_prev .* f
    x = mb.mul(x=ci, y=i)
    y = mb.mul(x=cs_prev, y=f)
    cs = mb.add(x=x, y=y)

    # cs = clip(cs, cell_clip)
    if op.cell_clip is not None:
        clip_val = op.cell_clip.val
        cs = mb.clip(x=cs, alpha=-clip_val, beta=clip_val)

    # o = sigmoid(cs * wco + o)
    if op.use_peephole.val:
        wco = op.weight_peep_o.val
        x = mb.mul(x=cs, y=wco)
        pre_o = mb.add(x=x, y=o)
    else:
        pre_o = o
    o = mb.sigmoid(x=pre_o)
    co = mb.tanh(x=cs)

    # h = co .* o
    h = mb.mul(x=co, y=o)

    return [i, cs, f, o, ci, co, h]


def _expand_tf_lstm_block_cell(op):
    if op.op_type != "tf_lstm_block_cell":
        raise ValueError()

    with op.enclosing_block as block:
        x = op.x  # [b, input_dim]
        h_prev = op.h_prev  # [b, hidden_dim]
        cs_prev = op.c_prev  # [b, hidden_dim]

        with mb.set_before_op(op):
            i, cs, f, o, ci, co, h = _lstm_cell_builder(op, x, h_prev, cs_prev)

        # Replace all outputs
        new_outputs = [i, cs, f, o, ci, co, h]
        for old_v, new_v in zip(op.outputs, new_outputs):
            block.replace_uses_of_var_after_op(
                anchor_op=op, old_var=old_v, new_var=new_v
            )
        block.remove_ops([op])


def _expand_tf_lstm_block(op):
    if op.op_type != "tf_lstm_block":
        raise ValueError()

    with op.enclosing_block as block:
        x = op.x  # [s, b, input_dim]
        h_prev = op.h_prev  # [b, hidden_dim]
        cs_prev = op.c_prev  # [b, hidden_dim]

        # cond and body function gor the while_loop
        def cond(i, cs_list, h_list):
            return mb.less(x=i, y=length)

        def body(i, cs_list, h_list):
            xi = mb.gather(x=x, indices=i, axis=0)
            h_prev = mb.gather(x=h_list, indices=i, axis=0)
            cs_prev = mb.gather(x=cs_list, indices=i, axis=0)

            ig, cs, fg, og, ci, co, h = _lstm_cell_builder(op, xi, h_prev, cs_prev)

            counter = mb.add(x=i, y=1)

            return (
                counter,
                mb.scatter(data=cs_list, indices=counter, updates=cs),
                mb.scatter(data=h_list, indices=counter, updates=h),
            )

        # Allocate two lists: cs & h
        x_shape = mb.shape(x=x, before_op=op)
        length = mb.slice_by_index(x=x_shape, begin=[0], end=[1], before_op=op)
        h_shape = mb.shape(x=h_prev, before_op=op)
        list_shape = mb.concat(values=[length, h_shape], axis=0, before_op=op)
        cs_list = mb.fill(shape=list_shape, before_op=op)
        h_list = mb.fill(shape=list_shape, before_op=op)

        # append initial state at index 0
        cs_prev = mb.expand_dims(x=cs_prev, axes=[0], before_op=op)
        cs_list = mb.concat(values=[cs_prev, cs_list], axis=0, before_op=op)
        h_prev = mb.expand_dims(x=h_prev, axes=[0], before_op=op)
        h_list = mb.concat(values=[h_prev, h_list], axis=0, before_op=op)

        _, cs_list, h_list = mb.while_loop(
            _cond=cond, _body=body, loop_vars=([0], cs_list, h_list), before_op=op
        )

        # strip initial state or element at index 0
        begin, end = [1, 0, 0], [0, 0, 0]
        begin_mask = [False, True, True]
        end_mask = [True, True, True]
        cs = mb.slice_by_index(
            x=cs_list,
            begin=begin,
            end=end,
            begin_mask=begin_mask,
            end_mask=end_mask,
            before_op=op,
        )
        h = mb.slice_by_index(
            x=h_list,
            begin=begin,
            end=end,
            begin_mask=begin_mask,
            end_mask=end_mask,
            before_op=op,
        )

        # Replace all outputs
        new_outputs = [cs, h]
        for old_v, new_v in zip(
            [ov for index, ov in enumerate(op.outputs) if index in [1, 6]], new_outputs
        ):
            block.replace_uses_of_var_after_op(
                anchor_op=op, old_var=old_v, new_var=new_v
            )
        block.remove_ops([op])
