"""Small CPU checks for release interfaces; no model assets or datasets required."""
import argparse
from pathlib import Path
import tempfile
import gzip
import json
import unittest
from unittest.mock import patch
from types import SimpleNamespace

import numpy as np
import torch
from torch import nn

from tools.config import parse_with_config, validate_paths
from tools.checkpoint import load_checkpoint, export_checkpoint
from eval.m4seg import greedy_mask_miou, image_size
from tools.outputs import image_mask_slots


class ConfigTests(unittest.TestCase):
    def test_hub_inference_defaults_preserve_backbone_training_initialization(self):
        from eval.tasks import parse_args as evaluation_args
        from training.trainer import parse_args as training_args
        from run import parse_args as inference_args
        with patch('sys.argv', ['run.py']):
            self.assertEqual(inference_args().merged_ckpt_path, 'PLAN-Lab/PRIMA')
        self.assertEqual(evaluation_args([]).merged_ckpt_path, 'PLAN-Lab/PRIMA')
        self.assertEqual(training_args([]).merged_ckpt_path, '')
        self.assertEqual(evaluation_args(['--merged_ckpt_path', '']).merged_ckpt_path, '')

    def test_cli_overrides_configuration(self):
        with tempfile.TemporaryDirectory() as directory:
            config = Path(directory) / 'options.yaml'
            config.write_text('batch_size: 4\n')
            parser = argparse.ArgumentParser()
            parser.add_argument('--batch_size', type=int, default=1)
            args = parse_with_config(parser, ['--config', str(config), '--batch_size', '2'])
            self.assertEqual(args.batch_size, 2)

    def test_placeholders_fail_before_model_loading(self):
        args = argparse.Namespace(version='/path/to/base_model')
        with self.assertRaisesRegex(ValueError, 'replace the placeholder'):
            validate_paths(args)

    def test_complete_hub_model_does_not_require_separate_assets(self):
        args = argparse.Namespace(merged_ckpt_path='PLAN-Lab/PRIMA',
                                  version='/path/to/base_model', ckpt_path='/path/to/prima_weights')
        validate_paths(args)
        args.dataset_dir = '/path/to/datasets'
        with self.assertRaisesRegex(ValueError, 'dataset_dir'):
            validate_paths(args, dataset=True)

    def test_complete_model_overlay_still_requires_an_existing_checkpoint(self):
        args = argparse.Namespace(merged_ckpt_path='PLAN-Lab/PRIMA',
                                  ckpt_path='/missing/prima_overlay.bin')
        with self.assertRaisesRegex(ValueError, 'ckpt_path'):
            validate_paths(args)

    def test_routing_uses_actual_seg_id_without_leading_whitespace(self):
        from tools.model_loader import setup_tokenizer_and_special_tokens, MODEL_DEFAULTS, INTERNAL_DEFAULTS

        class Tokenizer:
            unk_token = '<unk>'
            unk_token_id = 0

            def convert_tokens_to_ids(self, token):
                return 9

            def __call__(self, text, **kwargs):
                return SimpleNamespace(input_ids=[5, 9, 7, 8])

        args = SimpleNamespace(**MODEL_DEFAULTS, **INTERNAL_DEFAULTS)
        with patch('tools.model_loader.transformers.AutoTokenizer.from_pretrained', return_value=Tokenizer()):
            setup_tokenizer_and_special_tokens(args)
        self.assertEqual(args.seg_token_idx, 9)
        self.assertTrue(all(pattern == [9, 7, 8] for pattern in args.seg_image_tokens))


class PortableDataTests(unittest.TestCase):
    def test_four_digit_objects_keep_mask_order_and_five_digit_parts(self):
        from PIL import Image
        from pycocotools import mask as mask_utils
        from dataset.multi_image_datasets.M4seg_ds import M4SegDataset
        from model.llava import conversation
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for name in ('first.jpg', 'second.jpg'):
                Image.new('RGB', (6, 4), 'white').save(root / name)
            left = np.zeros((4, 6), dtype=np.uint8); left[:, :2] = 1
            right = np.zeros((4, 6), dtype=np.uint8); right[:, 4:] = 1
            def grounding(name, image, mask):
                rle = mask_utils.encode(np.asfortranarray(mask))
                rle['counts'] = rle['counts'].decode('ascii')
                return {'name': name, 'img': image, 'segmentation': rle}
            record = {'pair_id': 'example', 'dataset_name': 'coco_objects', 'img_ids': [1, 2],
                      'image_paths': ['first.jpg', 'second.jpg'], 'image_sizes': [[4, 6], [4, 6]],
                      'q': 'Compare book_2149 and chair_201.',
                      'a': 'The book_2149 and chair_201 have a leg_10102 nearby. Compare book_2149.',
                      'groundings': {'img_1': [grounding('leg', 1, left)],
                                     'img_2': [grounding('book', 2, left), grounding('chair', 2, right)]}}
            path = root / 'data.json'; path.write_text(json.dumps([record]))
            dataset = M4SegDataset(
                dataset_dir=root, 
                m4seg_annpath=path, mode='train')
            with patch.object(conversation, 'default_conversation', conversation.conv_templates['llava_v1']):
                item = dataset[0]
            text = item[4][0]
            self.assertIn('<p> book_149 </p> [SEG] (IMAGE2)', text)
            self.assertIn('<p> leg </p> [SEG] (IMAGE1)', text)
            self.assertEqual(text.count('<p> book_149 </p>'), 2)
            self.assertNotIn('book_2149', text)
            self.assertEqual(text.count('[SEG]'), 3)
            self.assertTrue(np.array_equal(item[5][2][0].numpy(), left))
            self.assertTrue(np.array_equal(item[5][2][1].numpy(), right))

            # Original annotations can repeat a grounding before a new entity.
            # Its stored answer span must select the new entity's own mask.
            record['a'] = 'The book_2149 repeats book_2149 before chair_201 and leg_10102.'
            record['groundings'] = {'img_1': [], 'img_2': []}
            import re
            for match in re.finditer(r'book_2149|chair_201|leg_10102', record['a']):
                name = match.group().split('_')[0]
                image = 1 if name == 'leg' else 2
                entry = grounding(name, image, right if name == 'chair' else left)
                entry.update(st_idx=match.start(), en_idx=match.end())
                record['groundings'][f'img_{image}'].append(entry)
            path.write_text(json.dumps([record]))
            dataset = M4SegDataset(
                dataset_dir=root, 
                m4seg_annpath=path, mode='train')
            with patch.object(conversation, 'default_conversation', conversation.conv_templates['llava_v1']):
                item = dataset[0]
            self.assertEqual(item[4][0].count('[SEG]'), 3)
            self.assertEqual(item[4][0].count('<p> book_149 </p>'), 2)
            self.assertTrue(np.array_equal(item[5][2][1].numpy(), right))

    def test_compressed_source_catalog_preserves_image_ids_without_mask_tables(self):
        from dataset.utils.coco import COCO
        from dataset.utils.utils import source_metadata_path
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            path = root / 'm4seg/metadata/source_catalogs/coco_objects/test.json.gz'
            path.parent.mkdir(parents=True)
            with gzip.open(path, 'wt') as handle:
                json.dump({'images': [{'id': 17, 'file_name': '000000000017.jpg',
                                       'height': 3, 'width': 4}],
                           'categories': [{'id': 2, 'name': 'object'}]}, handle)
            resolved = source_metadata_path(root, 'coco_objects', 'val', root / 'missing.json')
            api = COCO(resolved)
            self.assertEqual(api.getImgIds(), [17])
            self.assertEqual(api.loadImgs([17])[0]['width'], 4)

    def test_portable_gzip_annotations_need_no_source_annotation_files(self):
        from dataset.multi_image_datasets.M4seg_ds import M4SegDataset
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'annotations.json.gz'
            record = {'pair_id': 'sample', 'example_id': 'train-0000000',
                      'dataset_name': 'coco_objects', 'img_ids': [1],
                      'image_paths': ['coco_2017/train2017/example.jpg'], 'image_sizes': [[3, 4]],
                      'q': 'Describe the image.', 'a': 'An image.', 'groundings': {'img_1': []}}
            with gzip.open(path, 'wt') as handle:
                json.dump([record], handle)
            dataset = M4SegDataset(
                dataset_dir=directory, 
                m4seg_annpath=path, mode='test')
            self.assertTrue(dataset.portable_annotations)
            self.assertFalse(dataset.data2api)
            self.assertEqual(dataset.data_infos, [record])

    def test_portable_paths_reject_traversal(self):
        from dataset.multi_image_datasets.M4seg_ds import portable_image_metadata
        record = {'img_ids': [1], 'image_paths': ['../private.jpg'], 'image_sizes': [[3, 4]]}
        with self.assertRaisesRegex(ValueError, 'safe relative'):
            portable_image_metadata(record, '/path/to/images')

    def test_unique_example_ids_preserve_repeated_source_pair_ids(self):
        from eval.tasks import build_dataset
        from tools.model_loader import MODEL_DEFAULTS, INTERNAL_DEFAULTS
        args = SimpleNamespace(**MODEL_DEFAULTS, **INTERNAL_DEFAULTS,
                               dataset_dir='unused', m4seg_data='coco_objects', m4seg_annpath='unused')
        records = [{'pair_id': 'repeated', 'example_id': 'train-0000000'},
                   {'pair_id': 'repeated', 'example_id': 'train-0000001'}]
        dataset = SimpleNamespace(data_infos=records)
        with patch('eval.tasks.M4SegDataset', return_value=dataset):
            self.assertIs(build_dataset(args, mode='train'), dataset)

    def test_repeated_legacy_pair_ids_are_allowed_for_training_only(self):
        from eval.tasks import build_dataset
        from tools.model_loader import MODEL_DEFAULTS, INTERNAL_DEFAULTS
        args = SimpleNamespace(**MODEL_DEFAULTS, **INTERNAL_DEFAULTS,
                               dataset_dir='unused', m4seg_data='coco_objects', m4seg_annpath='unused')
        dataset = SimpleNamespace(data_infos=[{'pair_id': 'repeated'}, {'pair_id': 'repeated'}])
        with patch('eval.tasks.M4SegDataset', return_value=dataset):
            self.assertIs(build_dataset(args, mode='train'), dataset)
            with self.assertRaisesRegex(ValueError, 'unique'):
                build_dataset(args, mode='test')


def small_model():
    outer = nn.Module()
    outer.model = nn.Module()
    core = outer.model
    attention = nn.Module()
    attention.lora_A = nn.Linear(2, 1, bias=False)
    attention.lora_B = nn.Linear(1, 2, bias=False)
    core.layers = nn.ModuleList([attention])
    core.grounding_encoder = nn.Module()
    core.grounding_encoder.mask_decoder = nn.Linear(2, 2)
    core.query_tokens = nn.Parameter(torch.ones(1, 2))
    core.residual_query_tokens = nn.Parameter(torch.ones(1, 2))
    core.residual_extractor = nn.Linear(2, 2)
    core.mm_projector = nn.Linear(2, 2)
    core.text_hidden_fcs = nn.Linear(2, 2)
    core.visual_encoder = nn.Linear(2, 2)
    core.visual_encoder.requires_grad_(False)
    core.vision_tower = core.visual_encoder
    return outer


class CheckpointTests(unittest.TestCase):
    def test_export_omits_frozen_aliases_and_round_trips(self):
        source = small_model()
        destination = small_model()
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'weights.bin'
            export_checkpoint(source, path)
            state = torch.load(path, map_location='cpu')
            self.assertFalse(any('visual_encoder' in key or 'vision_tower' in key for key in state))
            load_checkpoint(destination, path)
            for key, value in state.items():
                self.assertTrue(torch.equal(destination.state_dict()[key], value), key)

    def test_incomplete_checkpoint_is_rejected(self):
        model = small_model()
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'weights.bin'
            export_checkpoint(model, path)
            state = torch.load(path, map_location='cpu')
            del state['model.mm_projector.weight']
            torch.save(state, path)
            with self.assertRaisesRegex(ValueError, 'Incomplete PRIMA checkpoint'):
                load_checkpoint(model, path)

    def test_peft_prefix_is_normalized(self):
        model = small_model()
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'weights.bin'
            export_checkpoint(model, path)
            state = torch.load(path, map_location='cpu')
            torch.save({'base_model.model.' + key: value for key, value in state.items()}, path)
            self.assertGreater(load_checkpoint(small_model(), path)['loaded_tensors'], 0)


class MetricTests(unittest.TestCase):
    def test_grounded_phrases_route_all_masks_and_normalize_labels(self):
        from eval.m4seg import extract_labels
        caption = ("<p> cabinet_17's drawer_2 </p> [SEG] (IMAGE1) [SEG] (IMAGE2) "
                   "<p> chair_12 </p> [SEG] (IMAGE2)")
        self.assertEqual(extract_labels(caption, 3),
                         [['drawer'], ['drawer', 'chair'], []])

    def test_caption_cleaning_removes_grounding_markup(self):
        from eval.m4seg import clean_caption, clean_reference_caption
        self.assertEqual(clean_caption('The <p> chair_12 </p> [SEG] (IMAGE2) is red.'),
                         'The chair  is red.')
        self.assertEqual(clean_reference_caption('The dining_chair_212 is red.'),
                         'The dining chair is red.')

    def test_direct_and_standalone_scoring_use_the_same_label_rules(self):
        from eval.m4seg import evaluate, score_predictions

        class Embedder:
            revision = 'test'

            def similarity(self, left, right):
                return float(left == right)

        grounding = {'name': 'chair_212', 'segmentation': [[0, 0, 4, 0, 4, 4, 0, 4]]}
        records = [{'pair_id': 'example', 'dataset_name': 'source', 'img_ids': [1, 2],
                    'a': 'A chair_101 and chair_201.',
                    'groundings': {'img_1': [grounding], 'img_2': [grounding]}}]
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            prediction = root / 'example.pth'
            torch.save({'output': ['<p> chair_101 </p> [SEG] (IMAGE1) [SEG] (IMAGE2)'],
                        'pred_masks': [None, torch.ones(1, 4, 4), torch.ones(1, 4, 4)],
                        'image_sizes': [[4, 4], [4, 4]]}, prediction)
            direct = evaluate(root, {'example': prediction}, records, {'example'},
                              embedder=Embedder())
            with patch('eval.m4seg.TextEmbedder', return_value=Embedder()):
                standalone = score_predictions([root], records, with_text=True)['results'][0]
            self.assertEqual(direct, standalone)
            self.assertEqual(direct['miou'], 1.0)
            self.assertEqual(direct['recall'], 1.0)
            self.assertEqual(direct['semantic_similarity'], 1.0)
            self.assertEqual(direct['iou_matched_semantic_similarity'], 1.0)

    def test_polygon_sizes_come_from_native_metadata(self):
        polygon = [{'segmentation': [[0, 0, 4, 0, 4, 3, 0, 3]]}]
        self.assertEqual(image_size({'image_sizes': [[3, 4]]}, 0, polygon), (3, 4))
        self.assertIsNone(image_size({}, 0, polygon))

    def test_legacy_sizes_can_be_read_from_source_images(self):
        from PIL import Image
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'image.png'
            Image.new('RGB', (7, 5)).save(path)
            self.assertEqual(image_size({'image_paths': [str(path)]}, 0, []), (5, 7))

    def test_matched_mask_denominator_is_preserved(self):
        mask = np.eye(4, dtype=bool)
        extra = np.zeros((4, 4), dtype=bool)
        self.assertEqual(greedy_mask_miou([mask, extra], [mask]), 1.0)
        self.assertEqual(greedy_mask_miou([], [mask]), 0.0)

    def test_explicit_grounding_slots_keep_source_order(self):
        first = torch.ones(1, 2, 2)
        second = torch.zeros(1, 3, 3)
        output = image_mask_slots([[None], [first], [second]], 2)
        self.assertIs(output[0], first)
        self.assertIs(output[1], second)
        self.assertEqual(image_mask_slots(None, 2), [None, None])

    def test_evaluation_can_generate_and_score_in_one_command(self):
        import json
        from eval import tasks

        records = [{'pair_id': 'example', 'dataset_name': 'source', 'img_ids': ['1'],
                    'groundings': {'img_1': [{'name': 'item',
                                             'segmentation': [[0, 0, 4, 0, 4, 4, 0, 4]]}]}}]

        class Dataset:
            data_infos = records

            def __len__(self):
                return len(self.data_infos)

        def generate(loader, model, tokenizer, args, output):
            torch.save({'output': ['<p> item </p> [SEG] (IMAGE1)'],
                        'pred_masks': [None, torch.ones(1, 4, 4)],
                        'image_sizes': [[4, 4]]}, output / 'example.pth')
            return {'samples': 1, 'written': 1, 'failures': {}}

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            annotation = root / 'annotations.json'
            annotation.write_text(json.dumps(records))
            checkpoint = root / 'checkpoint.bin'
            checkpoint.write_bytes(b'test')
            args = tasks.parse_args(['--merged_ckpt_path', '', '--ckpt_path', str(checkpoint), '--m4seg_annpath', str(annotation),
                                     '--output_save_path', str(root / 'outputs'), '--compute_metrics',
                                     '--no-metrics_with_text', '--no-metrics_with_meteor'])
            args.img_emb_len = 32
            with patch.object(tasks, 'validate_paths'), \
                    patch.object(tasks, 'load_bundle', return_value=(object(), object(), object())), \
                    patch.object(tasks, 'build_dataset', return_value=Dataset()), \
                    patch.object(tasks, 'inference', side_effect=generate):
                tasks.main(args)
            report = json.loads((root / 'outputs/metrics.json').read_text())
            self.assertEqual(report['results'][0]['miou'], 1.0)
            self.assertEqual(report['results'][0]['samples'], 1)

    def test_generation_records_native_image_dimensions(self):
        from eval import tasks

        class Loader:
            dataset = [0]

            def __iter__(self):
                yield {'sample_ids': ['example'], 'label_list': [(None,), (torch.zeros(3, 4),)],
                       'image_paths': [['image.png']], 'conversation_list': ['question']}

        model = SimpleNamespace(eval=lambda: None,
                                evaluate=lambda **kwargs: (torch.tensor([[1, 9, 2]]),
                                                            [[None], [torch.ones(1, 3, 4)]]))
        tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=2, pad_token='<pad>',
                                    batch_decode=lambda *args, **kwargs: ['answer'])
        args = SimpleNamespace(local_rank=0, resume=False, max_new_tokens=16)
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory)
            with patch.object(tasks, 'move_batch', side_effect=lambda batch, device: batch):
                summary = tasks.inference(Loader(), model, tokenizer, args, output)
            saved = torch.load(output / 'example.pth', map_location='cpu')
            self.assertEqual(saved['image_sizes'], [(3, 4)])
            self.assertEqual(summary['failures'], {})


class AttentionTests(unittest.TestCase):
    def test_local_causal_mask_preserves_backbone_semantics(self):
        from model.llava.model.language_model.prima_llama import PRIMAConfig, PRIMALlamaModel
        from transformers.models.llama.modeling_llama import LlamaModel

        config = PRIMAConfig(hidden_size=16, intermediate_size=32, num_attention_heads=4,
                             num_hidden_layers=1, vocab_size=32)
        config._attn_implementation = 'eager'
        model = PRIMALlamaModel(config)
        inputs = torch.randn(1, 3, 16)
        attention = torch.ones(1, 3)
        positions = torch.arange(3)
        expected = LlamaModel._update_causal_mask(model, attention, inputs, positions, None, False)
        actual = model._update_causal_mask(attention, inputs, positions, None, False)
        self.assertTrue(torch.equal(actual, expected))
        bf16 = model._update_causal_mask(attention, inputs.bfloat16(), positions, None, False)
        self.assertEqual(bf16.dtype, torch.bfloat16)
        self.assertEqual(bf16[0, 0, 0, 1].item(), torch.finfo(torch.bfloat16).min)


class TrainingTests(unittest.TestCase):
    def test_wrapped_training_keeps_frozen_encoders_in_eval_mode(self):
        from training.trainer import train_mode

        class Model(nn.Module):
            def __init__(self):
                super().__init__()
                self.core = nn.Module()
                for name in ('visual_encoder', 'ln_vision', 'Qformer'):
                    setattr(self.core, name, nn.Linear(2, 2))
                self.core.grounding_encoder = nn.Module()
                for name in ('image_encoder', 'prompt_encoder', 'mask_decoder'):
                    setattr(self.core.grounding_encoder, name, nn.Linear(2, 2))

            def get_model(self):
                return self.core

        wrapper = nn.Module()
        wrapper.module = Model()
        train_mode(wrapper)
        self.assertFalse(wrapper.module.core.visual_encoder.training)
        self.assertFalse(wrapper.module.core.Qformer.training)
        self.assertFalse(wrapper.module.core.grounding_encoder.prompt_encoder.training)
        self.assertTrue(wrapper.module.core.grounding_encoder.mask_decoder.training)


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