# Copyright 2024 Google LLC
#
# 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 kithara.model.maxtext.maxtext.MaxText.maxengine import MaxEngine
import jax
import max_utils

class MaxtextInferenceEngine(MaxEngine):
    """This is a patched version of MaxEngine
    
    Changes: 
      -  Added the `load_existing_params` function to allow 
      the engine to run with an exisitng model instead of initiating
      a new one. 
    """
    def load_existing_params(self, params, rng=None) -> "Params":

        if rng is None:
            rng = jax.random.PRNGKey(0)

        self.abstract_params = jax.tree_util.tree_map(
            lambda x: (
                jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding)
                if isinstance(x, jax.Array)
                else None
            ),
            params,
        )

        self.prefill_kv_cache_annotations = max_utils.get_prefill_kv_cache_annotations(
            self.model, self.config, rng, self._mesh
        )
        self.prefill_kv_cache_shardings = jax.tree_util.tree_map(
            lambda x: jax.sharding.NamedSharding(self._mesh, x),
            self.prefill_kv_cache_annotations,
        )

        if self.config.stack_prefill_result_cache:
            # Add extra axis for the axis generated by the stack.
            self.prefill_kv_cache_shardings = jax.tree_util.tree_map(
                lambda x: jax.sharding.NamedSharding(
                    self._mesh, jax.sharding.PartitionSpec(None, *x.spec)
                ),
                self.prefill_kv_cache_shardings,
            )
            self.prefill_kv_cache_shardings = self.prefill_kv_cache_shardings[
                "decoder"
            ]["layers_0"]

        self.kv_cache_annotations = max_utils.get_kv_cache_annotations(
            self.model, self.config, rng, self._mesh
        )
        self.kv_cache_shardings = jax.tree_util.tree_map(
            lambda x: jax.sharding.NamedSharding(self._mesh, x),
            self.kv_cache_annotations,
        )
        return params
