# Copyright (c) 2024 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 os
import unittest
from functools import partial
from tempfile import TemporaryDirectory
from types import SimpleNamespace

import paddle

from llm.utils.data import convert_example_for_reft
from paddlenlp.data import DataCollatorForSeq2Seq
from paddlenlp.datasets import load_dataset
from paddlenlp.peft.reft import (
    LoreftIntervention,
    LowRankRotateLayer,
    ReFTConfig,
    ReftDataCollator,
    ReFTModel,
    TinyIntervention,
    do_predict,
)
from paddlenlp.peft.reft.modeling_utils import (
    count_parameters,
    get_type_from_string,
    set_seed,
)
from paddlenlp.transformers import AutoModelForCausalLM, AutoTokenizer
from paddlenlp.trl import SFTTrainer


class TestReftDataCollator(unittest.TestCase):
    def test_call(self):
        model_name = "__internal_testing__/tiny-random-llama"
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            model_max_length=512,
            padding_side="right",
        )
        tokenizer.pad_token_id = tokenizer.eos_token_id
        model = AutoModelForCausalLM.from_pretrained(model_name)
        data_collator = DataCollatorForSeq2Seq(
            tokenizer=tokenizer, model=model, label_pad_token_id=-100, padding="longest"
        )
        reft_data_collator = ReftDataCollator(data_collator)
        instances = [
            {
                "input_ids": paddle.to_tensor([[1, 2, 3], [4, 5, 6]]),
                "intervention_locations": paddle.to_tensor([[0, 1, 0], [1, 0, 1]]),
            },
            {
                "input_ids": paddle.to_tensor([[7, 8, 9], [10, 11, 12]]),
                "intervention_locations": paddle.to_tensor([[1, 0, 1], [0, 1, 0]]),
            },
        ]

        batch_inputs = reft_data_collator(instances)

        self.assertIn("input_ids", batch_inputs)
        self.assertIn("intervention_locations", batch_inputs)
        self.assertIsInstance(batch_inputs["input_ids"], paddle.Tensor)
        self.assertIsInstance(batch_inputs["intervention_locations"], paddle.Tensor)


class TestBasicUtils(unittest.TestCase):
    def test_get_type_from_string(self):
        class_str = "paddlenlp.peft.reft.LoreftIntervention"
        cls = get_type_from_string(class_str)
        self.assertIsInstance(cls, type(LoreftIntervention))

    def test_set_seed(self):
        set_seed(42)
        set_seed(66)

    def test_count_param(self):
        model = AutoModelForCausalLM.from_pretrained("__internal_testing__/tiny-random-llama")
        count_parameters(model)


class TestReftConfig(unittest.TestCase):
    def test_reft_config(self):
        layers = [0, 1, 2]
        representations = [
            {
                "layer": l,
                "component": "block_output",
                "low_rank_dimension": 4,
                "intervention": LoreftIntervention(
                    embed_dim=768,
                    low_rank_dimension=4,
                    dropout=0.00,
                    dtype="float32",
                    act_fn="linear",
                    device="gpu",
                    add_bias=False,
                ),
            }
            for l in layers
        ]
        reft_config = ReFTConfig(representations=representations)
        reft_config.__str__()


class TestLoReftIntervention(unittest.TestCase):
    def setUp(self):
        self.kwargs = {
            "embed_dim": 64,
            "low_rank_dimension": 4,
            "dtype": paddle.float32,
            "dropout": 0.1,
            "act_fn": "linear",
        }

    def test_initialization(self):
        intervention = LoreftIntervention(**self.kwargs)
        self.assertIsInstance(intervention.rotate_layer, LowRankRotateLayer)
        self.assertIsInstance(intervention.learned_source, paddle.nn.Linear)
        self.assertEqual(intervention.dropout.p, self.kwargs["dropout"])

    def test_forward(self):
        base = paddle.randn([10, self.kwargs["embed_dim"]])
        intervention = LoreftIntervention(**self.kwargs)
        output = intervention.forward(base)
        self.assertEqual(output.shape, base.shape)
        self.assertEqual(output.dtype, self.kwargs["dtype"])

    def test_load_state_dict(self):
        model = LoreftIntervention(**self.kwargs)
        state_dict = {
            "learned_source.weight": paddle.randn([64, 4]),
            "learned_source.bias": paddle.zeros([4]),
            "rotate_layer.weight": paddle.randn([64, 4]),
        }
        model.load_state_dict(state_dict)
        self.assertTrue(paddle.allclose(model.learned_source.weight.data, state_dict["learned_source.weight"]))
        self.assertTrue(paddle.allclose(model.learned_source.bias.data, state_dict["learned_source.bias"]))
        self.assertTrue(
            paddle.allclose(
                model.rotate_layer.weight[:, : state_dict["rotate_layer.weight"].shape[-1]],
                state_dict["rotate_layer.weight"],
            )
        )


class TestTinyIntervention(unittest.TestCase):
    def setUp(self):
        self.kwargs = {
            "embed_dim": 768,
            "low_rank_dimension": 4,
            "dtype": paddle.float32,
            "dropout": 0.1,
            "act_fn": "relu",
        }

    def test_initialization(self):
        intervention = TinyIntervention(**self.kwargs)
        self.assertEqual(intervention.rank, self.kwargs["low_rank_dimension"])
        self.assertEqual(intervention.hidden_size, self.kwargs["embed_dim"])
        self.assertEqual(intervention.param_A.shape, [self.kwargs["embed_dim"], self.kwargs["low_rank_dimension"]])
        self.assertEqual(intervention.param_B.shape, [self.kwargs["low_rank_dimension"], self.kwargs["embed_dim"]])
        self.assertEqual(intervention.param_a.shape, [self.kwargs["low_rank_dimension"]])
        self.assertEqual(intervention.param_b.shape, [self.kwargs["embed_dim"]])

    def test_forward(self):
        base = paddle.randn([10, self.kwargs["embed_dim"]])
        intervention = TinyIntervention(**self.kwargs)
        output = intervention.forward(base)
        self.assertEqual(output.shape, base.shape)
        self.assertEqual(output.dtype, self.kwargs["dtype"])

    def test_load_state_dict(self):
        model = TinyIntervention(**self.kwargs)
        state_dict = {
            "param_A": paddle.randn([768, 4]),
            "param_B": paddle.randn([4, 768]),
            "param_a": paddle.randn([4]),
            "param_b": paddle.randn([768]),
        }
        model.load_state_dict(state_dict)
        self.assertTrue(paddle.allclose(model.param_A, state_dict["param_A"]))
        self.assertTrue(paddle.allclose(model.param_B, state_dict["param_B"]))
        self.assertTrue(paddle.allclose(model.param_a, state_dict["param_a"]))
        self.assertTrue(paddle.allclose(model.param_b, state_dict["param_b"]))


class TestReftModel(unittest.TestCase):
    def test_get_reft_model(self):
        model = AutoModelForCausalLM.from_pretrained("__internal_testing__/tiny-random-llama")
        layers = [0]
        representations = [
            {
                "layer": l,
                "component": "block_output",
                "low_rank_dimension": 4,
                "intervention": LoreftIntervention(
                    embed_dim=768,
                    low_rank_dimension=4,
                    dropout=0.00,
                    dtype="float32",
                    act_fn="linear",
                    device="gpu",
                    add_bias=False,
                ),
            }
            for l in layers
        ]
        reft_config = ReFTConfig(representations=representations)
        reft_model = ReFTModel(reft_config, model)
        reft_model.print_trainable_parameters()
        self.assertTrue(type(reft_model), ReFTModel)

    def test_reft_model_forward(self):
        model = AutoModelForCausalLM.from_pretrained("__internal_testing__/tiny-random-llama")

        layers = [0]
        representations = [
            {
                "layer": l,
                "component": "block_output",
                "low_rank_dimension": 4,
                "intervention": LoreftIntervention(
                    embed_dim=768,
                    low_rank_dimension=4,
                    dropout=0.00,
                    dtype="float32",
                    act_fn="linear",
                    device="gpu",
                    add_bias=False,
                ),
            }
            for l in layers
        ]
        reft_config = ReFTConfig(representations=representations)
        reft_model = ReFTModel(reft_config, model)
        reft_model.print_trainable_parameters()
        outputs = reft_model.model(**{"input_ids": paddle.randint(low=1, high=100, shape=(5, 10))})
        self.assertTrue(outputs[0].shape, [5, 10, 32000])


class TestReFTModelPredict(unittest.TestCase):
    def test_reft_model_predict(self):
        tokenizer = AutoTokenizer.from_pretrained("__internal_testing__/tiny-random-llama")
        tokenizer.pad_token_id = tokenizer.eos_token_id
        train_ds = load_dataset(
            "json",
            data_files=os.path.join("./tests/fixtures/llm/data", "train.json"),
            lazy=False,
        )[0]
        dev_ds = load_dataset(
            "json",
            data_files=os.path.join("./tests/fixtures/llm/data", "dev.json"),
            lazy=False,
        )[0]

        trans_func = partial(
            convert_example_for_reft,
            tokenizer=tokenizer,
            data_args=SimpleNamespace(**{"max_length": 64, "src_length": 32, "autoregressive": False}),
            positions="f7",
            num_interventions=1,
        )

        train_ds = train_ds.map(
            partial(
                trans_func,
                is_test=False,
                zero_padding=False,
                flash_mask=False,
            )
        )
        dev_ds = dev_ds.map(
            partial(
                trans_func,
                is_test=False,
                zero_padding=False,
                flash_mask=False,
            )
        )

        model = AutoModelForCausalLM.from_pretrained("__internal_testing__/tiny-random-llama")
        layers = [0]
        representations = [
            {
                "layer": l,
                "component": "block_output",
                "low_rank_dimension": 4,
                "intervention": LoreftIntervention(
                    embed_dim=768,
                    low_rank_dimension=4,
                    dropout=0.00,
                    dtype="float32",
                    act_fn="linear",
                    device="gpu",
                    add_bias=False,
                ),
            }
            for l in layers
        ]
        reft_config = ReFTConfig(representations=representations)
        reft_model = ReFTModel(reft_config, model)
        reft_model.disable_model_gradients()
        reft_model.model.train()
        reft_model.print_trainable_parameters()
        data_collator_fn = DataCollatorForSeq2Seq(
            tokenizer=tokenizer, model=model, label_pad_token_id=-100, padding="longest"
        )
        data_collator = ReftDataCollator(data_collator=data_collator_fn)
        trainer = SFTTrainer(
            model=reft_model,
            tokenizer=tokenizer,
            train_dataset=train_ds,
            data_collator=data_collator,
            eval_dataset=None,
            compute_metrics=None,
            gen_args=None,
            data_args=None,
            do_generation=False,
        )
        trainer.train()

        with TemporaryDirectory() as tempdir:
            reft_model.save_pretrained(tempdir)
            # 预测
            do_predict(
                intervenable=reft_model,
                tokenizer=tokenizer,
                eval_dataset=dev_ds,
                batch_size=1,
                predict_path=f"{tempdir}/pred_result.json",
            )


if __name__ == "__main__":
    unittest.main()
