# Copyright (c) 2023 PaddlePaddle Authors. 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.

import json
import os
import tempfile
import unittest

import paddle

from paddlenlp.transformers import (
    AutoConfig,
    BertModel,
    PretrainedConfig,
    PretrainedModel,
    register_base_model,
)
from paddlenlp.transformers.model_utils import load_sharded_checkpoint, shard_checkpoint
from paddlenlp.utils.env import (
    PADDLE_WEIGHTS_INDEX_NAME,
    PADDLE_WEIGHTS_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
)
from paddlenlp.utils.import_utils import is_paddle_cuda_available
from tests.testing_utils import require_package


class FakeConfig(PretrainedConfig):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)


class FakePretrainedModel(PretrainedModel):
    config_class = FakeConfig

    _keep_in_fp32_modules = ["norm."]


@register_base_model
class FakeModel(FakePretrainedModel):
    def __init__(self, config):
        super(FakeModel, self).__init__(config)
        self.linear = paddle.nn.Linear(2, 3)
        self.norm = paddle.nn.LayerNorm(2)


class TestFromPretrained(unittest.TestCase):
    def test_from_pretrained_low_cpu_mem_usage_functional(self):
        # test that we can use `from_pretrained(..., low_cpu_mem_usage=True)` with normal and
        # sharded models
        mnames = [
            "__internal_testing__/tiny-random-bert-sharded",
            "__internal_testing__/tiny-random-bert",
        ]
        for mname in mnames:
            m1 = BertModel.from_pretrained(mname, low_cpu_mem_usage=True)
            m2 = BertModel.from_pretrained(mname, low_cpu_mem_usage=False)
            for p1, p2 in zip(m1.parameters(), m2.parameters()):
                self.assertTrue(paddle.allclose(p1, p2))

    @unittest.skipIf(not is_paddle_cuda_available(), "some op is missing in cpu mode")
    def test_keep_in_fp32_modules(self):
        with tempfile.TemporaryDirectory() as tempdir:
            config = PretrainedConfig()
            model = FakeModel.from_config(config, dtype="float16")
            model.config = config
            model.save_pretrained(tempdir)

            # check model_state.pdparams
            state_dict = paddle.load(os.path.join(tempdir, "model_state.pdparams"))

            self.assertEqual(state_dict["linear.weight"].dtype, paddle.float16)
            self.assertEqual(state_dict["norm.weight"].dtype, paddle.float16)

            new_model = FakeModel.from_pretrained(tempdir)
            self.assertEqual(new_model.linear.weight.dtype, paddle.float16)
            self.assertEqual(new_model.norm.weight.dtype, paddle.float32)

    def test_load_sharded_checkpoint(self):
        config = AutoConfig.from_pretrained("__internal_testing__/bert-shard")
        model = BertModel.from_pretrained("__internal_testing__/bert-shard")

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir, max_shard_size="200kiB")
            model_load = BertModel.from_config(config)
            missing_keys, unexpected_keys = load_sharded_checkpoint(model_load, tmp_dir)

        self.assertEqual(missing_keys, [])
        self.assertEqual(unexpected_keys, [])
        for p1, p2 in zip(model.parameters(), model_load.parameters()):
            self.assertTrue(paddle.allclose(p1, p2))

    @unittest.skipIf(not is_paddle_cuda_available(), "some op is missing in cpu mode")
    def test_load_from_torch_dtyp_cast(self):
        pass

    @unittest.skipIf(not is_paddle_cuda_available(), "some op is missing in cpu mode")
    def test_load_dtype_cast(self):
        dtype_prefix_len = len("paddle.")

        def inner_convert_test(src_dtype, dst_dtype):
            str_src_dtype = str(src_dtype)[dtype_prefix_len:]
            str_dst_dtype = str(dst_dtype)[dtype_prefix_len:]

            config = AutoConfig.from_pretrained("__internal_testing__/tiny-random-bert")
            model = BertModel.from_config(config, dtype=str_src_dtype)

            with tempfile.TemporaryDirectory() as tmp_dir:
                model.save_pretrained(tmp_dir)
                new_model = BertModel.from_pretrained(tmp_dir, dtype=str_dst_dtype)

            for k, v in model.state_dict().items():
                if v.is_floating_point():
                    self.assertEqual(v.dtype, src_dtype)
            for k, v in new_model.state_dict().items():
                if v.is_floating_point():
                    self.assertEqual(v.dtype, dst_dtype)

        with self.subTest("paddle.float32 to paddle.float16"):
            inner_convert_test(paddle.float32, paddle.float16)
        with self.subTest("paddle.float32 to paddle.bfloat16"):
            inner_convert_test(paddle.float32, paddle.bfloat16)
        with self.subTest("paddle.float16 to paddle.float32"):
            inner_convert_test(paddle.float16, paddle.float32)
        with self.subTest("paddle.float16 to paddle.bfloat16"):
            inner_convert_test(paddle.float16, paddle.bfloat16)
        with self.subTest("paddle.bfloat16 to paddle.float32"):
            inner_convert_test(paddle.bfloat16, paddle.float32)
        with self.subTest("paddle.bfloat16 to paddle.float16"):
            inner_convert_test(paddle.bfloat16, paddle.float16)


class TestShardCheckpoint(unittest.TestCase):
    def test_shard_checkpoint(self):
        # This is the model we will use, total size 340,000 bytes.
        model = paddle.nn.Sequential(
            paddle.nn.Linear(100, 200, bias_attr=False),  # size 80,000
            paddle.nn.Linear(200, 200, bias_attr=False),  # size 160,000
            paddle.nn.Linear(200, 100, bias_attr=False),  # size 80,000
            paddle.nn.Linear(100, 50, bias_attr=False),  # size 20,000
        )
        state_dict = model.state_dict()

        with self.subTest("No shard when max size is bigger than model size"):
            shards, index = shard_checkpoint(state_dict)
            self.assertIsNone(index)
            self.assertDictEqual(shards, {PADDLE_WEIGHTS_NAME: state_dict})

        with self.subTest("Test sharding, no weights bigger than max size"):
            shards, index = shard_checkpoint(state_dict, max_shard_size="300kB")
            # Split is first two layers then last two.
            self.assertDictEqual(
                index,
                {
                    "metadata": {"total_size": 340000},
                    "weight_map": {
                        "0.weight": "model_state-00001-of-00002.pdparams",
                        "1.weight": "model_state-00001-of-00002.pdparams",
                        "2.weight": "model_state-00002-of-00002.pdparams",
                        "3.weight": "model_state-00002-of-00002.pdparams",
                    },
                },
            )

            shard1 = {"0.weight": state_dict["0.weight"], "1.weight": state_dict["1.weight"]}
            shard2 = {"2.weight": state_dict["2.weight"], "3.weight": state_dict["3.weight"]}
            self.assertDictEqual(
                shards, {"model_state-00001-of-00002.pdparams": shard1, "model_state-00002-of-00002.pdparams": shard2}
            )

        with self.subTest("Test sharding with weights bigger than max size"):
            shards, index = shard_checkpoint(state_dict, max_shard_size="100kB")
            # Split is first layer, second layer then last 2.
            self.assertDictEqual(
                index,
                {
                    "metadata": {"total_size": 340000},
                    "weight_map": {
                        "0.weight": "model_state-00001-of-00003.pdparams",
                        "1.weight": "model_state-00002-of-00003.pdparams",
                        "2.weight": "model_state-00003-of-00003.pdparams",
                        "3.weight": "model_state-00003-of-00003.pdparams",
                    },
                },
            )

            shard1 = {"0.weight": state_dict["0.weight"]}
            shard2 = {"1.weight": state_dict["1.weight"]}
            shard3 = {"2.weight": state_dict["2.weight"], "3.weight": state_dict["3.weight"]}
            self.assertDictEqual(
                shards,
                {
                    "model_state-00001-of-00003.pdparams": shard1,
                    "model_state-00002-of-00003.pdparams": shard2,
                    "model_state-00003-of-00003.pdparams": shard3,
                },
            )

    def test_checkpoint_sharding_local(self):
        model = BertModel.from_pretrained("__internal_testing__/bert-shard")

        with tempfile.TemporaryDirectory() as tmp_dir:
            # We use the same folder for various sizes to make sure a new save erases the old checkpoint.
            for max_size in ["50kB", "50kiB", "100kB", "100kiB", "200kB", "200kiB"]:
                model.save_pretrained(tmp_dir, max_shard_size=max_size)

                # Get each shard file and its size
                shard_to_size = {}
                for shard in os.listdir(tmp_dir):
                    if shard.endswith(".pdparams"):
                        shard_file = os.path.join(tmp_dir, shard)
                        shard_to_size[shard_file] = os.path.getsize(shard_file)

                index_file = os.path.join(tmp_dir, PADDLE_WEIGHTS_INDEX_NAME)
                # Check there is an index but no regular weight file
                self.assertTrue(os.path.isfile(index_file))
                self.assertFalse(os.path.isfile(os.path.join(tmp_dir, PADDLE_WEIGHTS_NAME)))

                # Check a file is bigger than max_size only when it has a single weight
                for shard_file, size in shard_to_size.items():
                    if max_size.endswith("kiB"):
                        max_size_int = int(max_size[:-3]) * 2**10
                    else:
                        max_size_int = int(max_size[:-2]) * 10**3
                    # Note: pickle adds some junk so the weight of the file can end up being slightly bigger than
                    # the size asked for (since we count parameters)
                    if size >= max_size_int + 50000:
                        state_dict = paddle.load(shard_file)
                        self.assertEqual(len(state_dict), 1)

                # Check the index and the shard files found match
                with open(index_file, "r", encoding="utf-8") as f:
                    index = json.loads(f.read())

                all_shards = set(index["weight_map"].values())
                shards_found = {f for f in os.listdir(tmp_dir) if f.endswith(".pdparams")}
                self.assertSetEqual(all_shards, shards_found)

                # Finally, check the model can be reloaded
                new_model = BertModel.from_pretrained(tmp_dir)
                for p1, p2 in zip(model.parameters(), new_model.parameters()):
                    self.assertTrue(paddle.allclose(p1, p2))

    def test_checkpoint_sharding_from_hub(self):
        model = BertModel.from_pretrained("__internal_testing__/tiny-random-bert-sharded")

        # the model above is the same as the model below, just a sharded version.
        ref_model = BertModel.from_pretrained("__internal_testing__/tiny-random-bert-no-sharded")
        for p1, p2 in zip(model.parameters(), ref_model.parameters()):
            self.assertTrue(paddle.allclose(p1, p2))

    def test_checkpoint_variant_local(self):
        model = BertModel.from_pretrained("__internal_testing__/tiny-random-bert")

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir, variant="v2")

            weights_name = ".".join(PADDLE_WEIGHTS_NAME.split(".")[:-1] + ["v2"] + ["pdparams"])

            weights_file = os.path.join(tmp_dir, weights_name)
            self.assertTrue(os.path.isfile(weights_file))
            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, PADDLE_WEIGHTS_NAME)))

            with self.assertRaises(EnvironmentError):
                _ = BertModel.from_pretrained(tmp_dir)

            new_model = BertModel.from_pretrained(tmp_dir, variant="v2")

        for p1, p2 in zip(model.parameters(), new_model.parameters()):
            self.assertTrue(paddle.allclose(p1, p2))

    def test_checkpoint_variant_local_sharded(self):
        model = BertModel.from_pretrained("__internal_testing__/tiny-random-bert")

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir, variant="v2", max_shard_size="50kB")

            weights_index_name = ".".join(PADDLE_WEIGHTS_INDEX_NAME.split(".")[:-1] + ["v2"] + ["json"])
            weights_index_file = os.path.join(tmp_dir, weights_index_name)
            self.assertTrue(os.path.isfile(weights_index_file))
            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, PADDLE_WEIGHTS_INDEX_NAME)))

            for i in range(1, 6):
                weights_name = ".".join(PADDLE_WEIGHTS_NAME.split(".")[:-1] + [f"v2-0000{i}-of-00005"] + ["pdparams"])
                weights_name_file = os.path.join(tmp_dir, weights_name)
                self.assertTrue(os.path.isfile(weights_name_file))

            with self.assertRaises(EnvironmentError):
                _ = BertModel.from_pretrained(tmp_dir)

            new_model = BertModel.from_pretrained(tmp_dir, variant="v2")

        for p1, p2 in zip(model.parameters(), new_model.parameters()):
            self.assertTrue(paddle.allclose(p1, p2))

    @require_package("safetensors")
    def test_checkpoint_variant_local_safe(self):
        model = BertModel.from_pretrained("__internal_testing__/tiny-random-bert")

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir, variant="v2", safe_serialization=True)

            weights_name = ".".join(SAFE_WEIGHTS_NAME.split(".")[:-1] + ["v2"] + ["safetensors"])

            weights_file = os.path.join(tmp_dir, weights_name)

            self.assertTrue(os.path.isfile(weights_file))
            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))

            with self.assertRaises(EnvironmentError):
                _ = BertModel.from_pretrained(tmp_dir)

            new_model = BertModel.from_pretrained(tmp_dir, variant="v2")

        for p1, p2 in zip(model.parameters(), new_model.parameters()):
            self.assertTrue(paddle.allclose(p1, p2))

    @require_package("safetensors")
    def test_checkpoint_variant_local_sharded_safe(self):
        model = BertModel.from_pretrained("__internal_testing__/tiny-random-bert")

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir, variant="v2", max_shard_size="50kB", safe_serialization=True)

            weights_index_name = ".".join(SAFE_WEIGHTS_INDEX_NAME.split(".")[:-1] + ["v2"] + ["json"])
            weights_index_file = os.path.join(tmp_dir, weights_index_name)
            self.assertTrue(os.path.isfile(weights_index_file))
            self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)))

            for i in range(1, 6):
                weights_name = ".".join(SAFE_WEIGHTS_NAME.split(".")[:-1] + [f"v2-0000{i}-of-00005"] + ["safetensors"])
                weights_name_file = os.path.join(tmp_dir, weights_name)
                self.assertTrue(os.path.isfile(weights_name_file))

            with self.assertRaises(EnvironmentError):
                _ = BertModel.from_pretrained(tmp_dir)

            new_model = BertModel.from_pretrained(tmp_dir, variant="v2")

        for p1, p2 in zip(model.parameters(), new_model.parameters()):
            self.assertTrue(paddle.allclose(p1, p2))

    def test_checkpoint_variant_hub(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            with self.assertRaises(EnvironmentError):
                _ = BertModel.from_pretrained("__internal_testing__/tiny-random-bert-variant", cache_dir=tmp_dir)

            model = BertModel.from_pretrained(
                "__internal_testing__/tiny-random-bert-variant", cache_dir=tmp_dir, variant="v2"
            )
        self.assertIsNotNone(model)

    def test_checkpoint_variant_hub_sharded(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            with self.assertRaises(EnvironmentError):
                _ = BertModel.from_pretrained(
                    "__internal_testing__/tiny-random-bert-variant-sharded", cache_dir=tmp_dir
                )
            model = BertModel.from_pretrained(
                "__internal_testing__/tiny-random-bert-variant-sharded", cache_dir=tmp_dir, variant="v2"
            )
        self.assertIsNotNone(model)

    def test_checkpoint_variant_save_load(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            model = BertModel.from_pretrained(
                "__internal_testing__/tiny-random-bert-variant", cache_dir=tmp_dir, variant="v2"
            )
            weights_name = ".".join(PADDLE_WEIGHTS_NAME.split(".")[:-1] + ["v2"] + ["pdparams"])

            model.save_pretrained(tmp_dir, variant="v2")
            # saving will create a variant checkpoint
            self.assertTrue(os.path.isfile(os.path.join(tmp_dir, weights_name)))

            model.save_pretrained(tmp_dir)
            # saving shouldn't delete variant checkpoints
            weights_name = ".".join(PADDLE_WEIGHTS_NAME.split(".")[:-1] + ["v2"] + ["pdparams"])
            self.assertTrue(os.path.isfile(os.path.join(tmp_dir, weights_name)))

            # there should be a normal checkpoint
            self.assertTrue(os.path.isfile(os.path.join(tmp_dir, PADDLE_WEIGHTS_NAME)))

        self.assertIsNotNone(model)
