import json
import math
import tempfile
import unittest
from pathlib import Path

import torch
import torch.nn.functional as F

from qprefer_reward.training.data import (
    QPreferPairCollator,
    load_manifest,
    parse_path_prefix_maps,
    resolve_media_path,
)
from qprefer_reward.training.loss import rao_kupper_loss


class TrainingLossTest(unittest.TestCase):
    def test_rao_kupper_decisive_tie_and_mask(self):
        rewards_a = torch.tensor([[1.0, 0.0, 7.0]])
        rewards_b = torch.tensor([[0.0, 0.0, -3.0]])
        labels = torch.tensor([[1, 0, 22]])
        observed = rao_kupper_loss(rewards_a, rewards_b, labels, k=5.0)

        log_k = math.log(5.0)
        decisive = -F.logsigmoid(torch.tensor(1.0 - log_k))
        tie = (
            -F.logsigmoid(torch.tensor(-log_k)) * 2
            - math.log(5.0**2 - 1.0)
        )
        expected = (decisive + tie) / 3
        self.assertTrue(torch.allclose(observed, expected))

    def test_unit_sample_weights_preserve_original_scale(self):
        generator = torch.Generator().manual_seed(7)
        rewards_a = torch.randn(4, 3, generator=generator)
        rewards_b = torch.randn(4, 3, generator=generator)
        labels = torch.tensor([[1, 22, 0], [-1, 22, 1], [0, 22, -1], [1, 22, 1]])
        unweighted = rao_kupper_loss(rewards_a, rewards_b, labels)
        weighted = rao_kupper_loss(
            rewards_a,
            rewards_b,
            labels,
            sample_weight=torch.ones(4),
        )
        self.assertTrue(torch.allclose(unweighted, weighted))


class TrainingManifestTest(unittest.TestCase):
    def test_relative_media_paths_and_summary(self):
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            media = root / "media"
            media.mkdir()
            for name in ("a.mp4", "b.mp4", "ref.jpg"):
                (media / name).touch()
            manifest = root / "pairs.json"
            manifest.write_text(
                json.dumps(
                    [
                        {
                            "pair_id": "pair-1",
                            "task": "i2v",
                            "bucket": "near",
                            "source_version": "unit",
                            "prompt": "A test prompt",
                            "path_A": "a.mp4",
                            "path_B": "b.mp4",
                            "ref_image": "ref.jpg",
                            "chosen_label": [1, 22, 0],
                        }
                    ]
                )
            )
            rows, summary = load_manifest(manifest, media_root=media, check_media=True)
            self.assertEqual(rows[0]["path_A"], str((media / "a.mp4").resolve()))
            self.assertEqual(summary["rows"], 1)
            self.assertEqual(summary["labels"]["MQ"], {"22": 1})

    def test_training_prompt_is_the_inference_contract(self):
        prompt = QPreferPairCollator.reward_prompt("A red fox runs.")
        self.assertIn("Visual Quality: <|VQ_reward|>", prompt)
        self.assertIn("Motion Quality: <|MQ_reward|>", prompt)
        self.assertIn("Text/Image Alignment: <|TA_reward|>", prompt)
        self.assertTrue(prompt.endswith("Textual prompt - A red fox runs.\n"))

    def test_absolute_paths_can_be_relocated_without_editing_manifest(self):
        mappings = parse_path_prefix_maps(["/old/project=/new/dataset"])
        observed = resolve_media_path(
            "/old/project/videos/a.mp4",
            Path("/tmp/manifest.json"),
            None,
            mappings,
        )
        self.assertEqual(observed, "/new/dataset/videos/a.mp4")


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