# Copyright (c) 2023 Graphcore Ltd.
# 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.

import math
import numpy as np
from poptransformer import ops
from poptransformer.utils import build_sharded_weight
from .base_layer import BaseLayer


class BaseEmbedding(BaseLayer):

    def __init__(self, context, name, vocab_size, embed_size):
        super().__init__(context, name)
        self.vocab_size = vocab_size
        self.embed_size = embed_size
        self.collect_bind_layer_weights()

    def collect_bind_layer_weights(self):
        weight_key = '.'.join([self.context, 'weight'])
        weight_np = self.get_param_from_state_dict(weight_key, [self.vocab_size, self.embed_size])
        self.weight_id = self.add_initialized_input_tensor(weight_np, weight_key)

    def __call__(self, graph, input_ids, sequence_length):
        with graph.nameScope(self.context):
            return ops.gather(graph, self.weight_id, input_ids)


class TPEmbedding(BaseEmbedding):

    def collect_bind_layer_weights(self):
        self.vs_setting = {'vs_type': 'consecutive', 'group_size': 1}
        self.vocab_per_ipu = math.ceil(self.vocab_size / self.num_replicas)

        weight_key = '.'.join([self.context, 'weight'])
        weight_np = self.get_param_from_state_dict(weight_key, [self.vocab_size, self.embed_size])
        weight_np = build_sharded_weight(weight_np, self.num_replicas, self.vocab_size, self.embed_size)
        self.weight_id = self.add_initialized_input_tensor(weight_np, weight_key, **self.vs_setting)

        index_offset_np = np.expand_dims(np.arange(self.num_replicas, dtype=np.int32), [1, 2]) * self.vocab_per_ipu
        self.index_offset = self.add_initialized_input_tensor(index_offset_np, 'index_offset', **self.vs_setting)

    def __call__(self, graph, input_ids, sequence_length):
        with graph.nameScope(self.context):
            replicated_input_ids = ops.sub(graph, input_ids, self.index_offset)
            cond1 = ops.less(graph, replicated_input_ids, ops.constant(graph, np.array([0], dtype=np.int32)))
            cond2 = ops.greater(
                graph, replicated_input_ids, ops.constant(graph, np.array([self.vocab_per_ipu], dtype=np.int32)))
            cond = ops.bitwiseor(graph, ops.cast(graph, cond1, 'INT32'), ops.cast(graph, cond2, 'INT32'))
            cond = ops.cast(graph, cond, 'BOOL')
            fill_value_np = self.vocab_per_ipu * np.ones(
                (self.num_replicas, self.batch_size, sequence_length), dtype=np.int32)
            fill_value = self.add_initialized_input_tensor(fill_value_np, 'fill_value', **self.vs_setting)
            updated_replicated_input_ids = ops.where(graph, cond, fill_value, replicated_input_ids)
            output = ops.gather(graph, self.weight_id, updated_replicated_input_ids)
            return output


class Embedding(TPEmbedding, BaseEmbedding):
    layer_class_map = {
        'tp': TPEmbedding,
        'shard': BaseEmbedding}

    def __init__(self, context, name, vocab_size, embed_size):
        model_type = self.model_type
        self.layer_class = self.layer_class_map.get(model_type, None)
        if not self.layer_class:
            raise ValueError(f"Invalid model_type {model_type}, options: {self.layer_class_map.keys()}")
        self.logger.debug(f'initializing model type: {self.layer_class.__name__}')
        super().__init__(context, name, vocab_size, embed_size)

    def __call__(self, graph, input_ids, sequence_length):
        return self.layer_class.__call__(self, graph, input_ids, sequence_length)

    def collect_bind_layer_weights(self):
        return self.layer_class.collect_bind_layer_weights(self)