"""Check training-only head initialization without weakening predictor loading."""
import importlib.util
from pathlib import Path
import tempfile
import unittest

ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location('regression_loader', ROOT / 'validate_release.py')
loader = importlib.util.module_from_spec(spec)
spec.loader.exec_module(loader)


class RegressionBaseTests(unittest.TestCase):
    def test_shared_default(self):
        self.assertEqual(loader.REGRESSION_BASE, 'RNASeek-regression-pretrained')
        self.assertTrue(loader.local_path(loader.REGRESSION_BASE).is_dir())
        self.assertTrue(loader.local_path('RNASeek-lm-pretrained').is_dir())
        for task in loader.DEFAULTS.values():
            self.assertNotEqual(task['checkpoint'], loader.REGRESSION_BASE)

    def test_pretrained_and_predictor_loading(self):
        import torch
        from safetensors.torch import save_file
        from tokenizers import Tokenizer, models
        from transformers import Qwen2Config, Qwen2ForCausalLM, PreTrainedTokenizerFast
        with tempfile.TemporaryDirectory() as temporary:
            base = Path(temporary) / 'base'
            config = Qwen2Config(vocab_size=8, hidden_size=8, intermediate_size=16,
                num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=1,
                max_position_embeddings=32, pad_token_id=0, bos_token_id=1, eos_token_id=2)
            original = Qwen2ForCausalLM(config)
            original.save_pretrained(base)
            tokenizer = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(models.WordLevel(
                {'<pad>': 0, '<s>': 1, '</s>': 2, '<unk>': 3, 'A': 4}, unk_token='<unk>')),
                pad_token='<pad>', bos_token='<s>', eos_token='</s>', unk_token='<unk>')
            tokenizer.save_pretrained(base)
            with self.assertRaisesRegex(ValueError, 'no trained regression head'):
                loader.load_regression_checkpoint(base)
            backbone, head, _ = loader.load_regression_checkpoint(base, initialize_head=True)
            for key, value in original.model.state_dict().items():
                torch.testing.assert_close(backbone.state_dict()[key], value, rtol=0, atol=0)
            head.net(torch.randn(2, 8)).sum().backward()
            self.assertIsNotNone(head.net[-1].weight.grad)
            predictor = Path(temporary) / 'predictor'
            predictor.mkdir()
            state = {'backbone.' + k: v for k, v in backbone.state_dict().items()}
            state.update({'regression_head.' + k: v for k, v in head.state_dict().items()})
            save_file(state, str(predictor / 'model.safetensors'))
            config.save_pretrained(predictor)
            tokenizer.save_pretrained(predictor)
            for training in (False, True):
                _, retained_head, _ = loader.load_regression_checkpoint(predictor, initialize_head=training)
                for key, value in head.state_dict().items():
                    torch.testing.assert_close(retained_head.state_dict()[key], value, rtol=0, atol=0)