import os
import shutil
import tempfile
import unittest

import torch
from modelscope import AutoModel, Preprocessor
from peft.utils import WEIGHTS_NAME
from transformers import PreTrainedModel

from swift import LoRAConfig, Swift
from swift.tuners import NEFTuneConfig


class TestNEFT(unittest.TestCase):

    def setUp(self):
        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
        self.tmp_dir = tempfile.TemporaryDirectory().name
        if not os.path.exists(self.tmp_dir):
            os.makedirs(self.tmp_dir)

    def tearDown(self):
        shutil.rmtree(self.tmp_dir)
        super().tearDown()

    def test_neft(self):
        model = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
        preprocessor = Preprocessor.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
        inputs = preprocessor('how are you')
        config = NEFTuneConfig()

        t1 = model.embeddings.word_embeddings(inputs['input_ids'])
        model = Swift.prepare_model(model, config)
        model.train()
        t2 = model.embeddings.word_embeddings(inputs['input_ids'])
        model.deactivate_adapter('default')
        t3 = model.embeddings.word_embeddings(inputs['input_ids'])
        self.assertTrue(torch.allclose(t1, t3))
        self.assertFalse(torch.allclose(t1, t2))
        model.save_pretrained(self.tmp_dir)
        bin_file = os.path.join(self.tmp_dir, 'pytorch_model.bin')
        self.assertTrue(os.path.isfile(bin_file))
        model2 = AutoModel.from_pretrained(self.tmp_dir)

        state_dict = model.state_dict()
        state_dict2 = model2.state_dict()
        self.assertTrue(len(state_dict) > 0)
        for key in state_dict:
            self.assertTrue(key in state_dict2)
            self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu()))

        shutil.rmtree(self.tmp_dir)
        PreTrainedModel.origin_save_pretrained = PreTrainedModel.save_pretrained
        delattr(PreTrainedModel, 'save_pretrained')
        model.save_pretrained(self.tmp_dir)
        bin_file = os.path.join(self.tmp_dir, WEIGHTS_NAME)
        self.assertTrue(os.path.isfile(bin_file))
        model_new = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
        model_new_2 = Swift.from_pretrained(model_new, self.tmp_dir)

        state_dict = model.state_dict()
        state_dict2 = model_new_2.state_dict()
        self.assertTrue(len(state_dict) > 0)
        for key in state_dict:
            self.assertTrue(key in state_dict2)
            self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu()))
        PreTrainedModel.save_pretrained = PreTrainedModel.origin_save_pretrained

    def test_neft_lora(self):
        model = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
        preprocessor = Preprocessor.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
        inputs = preprocessor('how are you')
        config = NEFTuneConfig()
        config2 = LoRAConfig(target_modules=['query', 'key', 'value'])

        t1 = model.embeddings.word_embeddings(inputs['input_ids'])
        model = Swift.prepare_model(model, {'c1': config, 'c2': config2})
        model.train()
        t2 = model.embeddings.word_embeddings(inputs['input_ids'])
        model.deactivate_adapter('c1')
        t3 = model.embeddings.word_embeddings(inputs['input_ids'])
        self.assertTrue(torch.allclose(t1, t3))
        self.assertFalse(torch.allclose(t1, t2))
        model.save_pretrained(self.tmp_dir)
        bin_file = os.path.join(self.tmp_dir, 'c2', WEIGHTS_NAME)
        self.assertTrue(os.path.isfile(bin_file))
        bin_file = os.path.join(self.tmp_dir, 'c1', WEIGHTS_NAME)
        self.assertTrue(not os.path.isfile(bin_file))
        model_new = AutoModel.from_pretrained('AI-ModelScope/bert-base-uncased')
        t1 = model_new.embeddings.word_embeddings(inputs['input_ids'])
        model_new = Swift.from_pretrained(model_new, self.tmp_dir)
        model_new.train()
        t2 = model_new.embeddings.word_embeddings(inputs['input_ids'])
        model_new.eval()
        t4 = model_new.embeddings.word_embeddings(inputs['input_ids'])
        model_new.train()
        model_new.deactivate_adapter('c1')
        t3 = model_new.embeddings.word_embeddings(inputs['input_ids'])
        self.assertTrue(torch.allclose(t1, t3))
        self.assertTrue(torch.allclose(t1, t4))
        self.assertFalse(torch.allclose(t1, t2))

        state_dict = model.state_dict()
        state_dict2 = model_new.state_dict()
        self.assertTrue(len(state_dict) > 0 and all(['lora' in key for key in state_dict.keys()]))
        for key in state_dict:
            self.assertTrue(key in state_dict2)
            self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu()))