import os.path
import shutil
import tempfile
import unittest

import torch
from modelscope import Model

from swift import LoRAConfig, Swift
from swift.tuners.utils import ModulesToSaveWrapper


class TestExtraStateDict(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_swift_extra_state_dict(self):
        model = Model.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
        lora_config = LoRAConfig(target_modules=['query', 'key', 'value'])
        model = Swift.prepare_model(model, lora_config, extra_state_keys=['classifier.*'])
        model.save_pretrained(self.tmp_dir)
        self.assertTrue(os.path.isfile(os.path.join(self.tmp_dir, 'extra_states', 'adapter_model.bin')))
        state_dict = torch.load(os.path.join(self.tmp_dir, 'extra_states', 'adapter_model.bin'))
        self.assertTrue(any('classifier' in key for key in state_dict))
        state_dict['classifier.weight'] = torch.ones_like(state_dict['classifier.weight']) * 2.0
        with open(os.path.join(self.tmp_dir, 'extra_states', 'adapter_model.bin'), 'wb') as f:
            torch.save(state_dict, f)
        model = Model.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
        model = Swift.from_pretrained(model, self.tmp_dir, inference_mode=False)
        names = [name for name, value in model.named_parameters() if value.requires_grad]
        self.assertTrue(any('classifier' in name for name in names))
        self.assertTrue(torch.allclose(state_dict['classifier.weight'], model.base_model.classifier.weight))

    def test_swift_modules_to_save(self):
        model = Model.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
        lora_config = LoRAConfig(target_modules=['query', 'key', 'value'], modules_to_save=['classifier'])
        lora_config2 = LoRAConfig(target_modules=['query', 'key', 'value'], modules_to_save=['classifier'])
        model = Swift.prepare_model(model, {'lora1': lora_config, 'lora2': lora_config2})
        model.set_active_adapters('lora1')
        model.set_active_adapters('lora2')
        self.assertTrue(isinstance(model.classifier, ModulesToSaveWrapper))
        self.assertTrue(model.classifier.active_adapter == 'lora2')
        model.save_pretrained(self.tmp_dir)
        state_dict = torch.load(os.path.join(self.tmp_dir, 'lora2', 'adapter_model.bin'))
        self.assertTrue(any('classifier' in key for key in state_dict))
        state_dict['classifier.weight'] = torch.ones_like(state_dict['classifier.weight']) * 2.0
        with open(os.path.join(self.tmp_dir, 'lora2', 'adapter_model.bin'), 'wb') as f:
            torch.save(state_dict, f)
        model = Model.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
        model = Swift.from_pretrained(model, self.tmp_dir, adapter_name='lora2')
        names = [name for name, value in model.named_parameters() if value.requires_grad]
        self.assertTrue(any('classifier' in name for name in names))
        self.assertTrue(
            torch.allclose(state_dict['classifier.weight'],
                           model.base_model.classifier.modules_to_save['lora2'].weight))