"""Focused tests for EasyCom loss/output-space ablations."""

from __future__ import annotations

import math
import unittest

import torch

from easycom_head import (
    TurnPredictionHead,
    activity_targets_to_four_states,
    checkpoint_loss_name,
    four_state_multihorizon_loss,
    independent_bce_multihorizon_loss,
    marginal_probabilities_from_logits,
    output_classes_for_loss,
)


class ObjectiveTests(unittest.TestCase):
    def setUp(self) -> None:
        self.targets = torch.tensor(
            [
                [1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 0.0, 1.0],
                [0.0, 1.0, 1.0],
            ]
        )

    def test_activity_target_mapping(self) -> None:
        actual = activity_targets_to_four_states(self.targets)
        self.assertTrue(torch.equal(actual, torch.tensor([0, 1, 2, 3])))

    def test_invalid_activity_target_is_rejected(self) -> None:
        with self.assertRaisesRegex(ValueError, "must be one of"):
            activity_targets_to_four_states(torch.tensor([[0.0, 0.0, 0.0]]))

    def test_independent_bce_uses_standard_mean(self) -> None:
        logits = torch.zeros_like(self.targets)
        loss = independent_bce_multihorizon_loss(logits, self.targets)
        self.assertAlmostEqual(float(loss), math.log(2.0), places=6)

    def test_four_state_loss_accepts_three_label_targets(self) -> None:
        logits = torch.zeros(4, 4)
        loss = four_state_multihorizon_loss(logits, self.targets)
        self.assertAlmostEqual(float(loss), math.log(4.0), places=6)

    def test_four_state_probabilities_are_marginalized(self) -> None:
        state_probabilities = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
        logits = state_probabilities.log()
        actual = marginal_probabilities_from_logits(logits, "four-state-ce")
        expected = torch.tensor([[0.1, 0.6, 0.7]])
        self.assertTrue(torch.allclose(actual, expected, atol=1e-6))
        self.assertGreater(float(actual.sum()), 1.0)

    def test_bce_probabilities_use_sigmoid(self) -> None:
        actual = marginal_probabilities_from_logits(
            torch.zeros(2, 3),
            "independent-bce",
        )
        self.assertTrue(torch.equal(actual, torch.full((2, 3), 0.5)))

    def test_four_state_head_shape(self) -> None:
        head = TurnPredictionHead(
            hidden_size=8,
            num_horizons=2,
            num_classes=output_classes_for_loss("four-state-ce"),
            dropout=0.0,
        )
        self.assertEqual(tuple(head(torch.randn(3, 8)).shape), (3, 2, 4))

    def test_all_new_objectives_backpropagate_to_marginals(self) -> None:
        states = torch.randn(4, 8)
        future_targets = self.targets[:, None, :].repeat(1, 2, 1)
        for loss_name, loss_function in (
            ("independent-bce", independent_bce_multihorizon_loss),
            ("four-state-ce", four_state_multihorizon_loss),
        ):
            with self.subTest(loss=loss_name):
                head = TurnPredictionHead(
                    hidden_size=8,
                    num_horizons=2,
                    num_classes=output_classes_for_loss(loss_name),
                    dropout=0.0,
                )
                logits = head(states)
                loss = loss_function(logits, future_targets)
                loss.backward()
                scores = marginal_probabilities_from_logits(logits, loss_name)
                self.assertEqual(tuple(scores.shape), (4, 2, 3))
                self.assertTrue(torch.isfinite(scores).all())
                self.assertIsNotNone(head.net[-1].weight.grad)

    def test_legacy_checkpoint_defaults_to_lstr(self) -> None:
        self.assertEqual(checkpoint_loss_name({}), "lstr-mce")

    def test_checkpoint_rejects_loss_output_mismatch(self) -> None:
        payload = {
            "training_config": {"loss": "four-state-ce"},
            "head_config": {"num_classes": 3},
        }
        with self.assertRaisesRegex(ValueError, "requires 4"):
            checkpoint_loss_name(payload)


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