import os
import sys
import unittest
import torch
import torch.nn as nn

# Ensure python directory is in path
sys.path.append(
    os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "python")
)

from parasitic_qlora import ExpertAdapter, LoRAMatrices
from adapter_moe_router import ExpertAdapterRouter


class SimpleModel(nn.Module):  # type: ignore[misc]
    def __init__(self) -> None:
        super().__init__()
        self.fc1 = nn.Linear(32, 16, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc1(x)


class TestAdapterMoERouter(unittest.TestCase):
    def setUp(self) -> None:
        self.model = SimpleModel()

        # Create dummy expert adapters with domain tags
        # Adapter 1: Statute Recall (has fc1.weight adapter)
        self.adapter1 = ExpertAdapter(
            adapter_id="adapter_statute",
            step=1,
            domain_tags=["statute_recall"],
            layers={
                "fc1.weight": LoRAMatrices(
                    layer_name="fc1.weight",
                    lora_B=torch.ones(16, 2) * 0.1,  # d x r = 16 x 2
                    lora_A=torch.ones(2, 32) * 0.1,  # r x k = 2 x 32
                    rank=2,
                    explained_variance=1.0,
                    singular_values=torch.tensor([1.0, 1.0]),
                    original_shape=(16, 32),
                )
            },
        )

        # Adapter 2: Logic (has fc1.weight adapter)
        self.adapter2 = ExpertAdapter(
            adapter_id="adapter_logic",
            step=1,
            domain_tags=["logic_reasoning"],
            layers={
                "fc1.weight": LoRAMatrices(
                    layer_name="fc1.weight",
                    lora_B=torch.ones(16, 2) * 0.2,
                    lora_A=torch.ones(2, 32) * 0.2,
                    rank=2,
                    explained_variance=1.0,
                    singular_values=torch.tensor([1.0, 1.0]),
                    original_shape=(16, 32),
                )
            },
        )

        self.library = [self.adapter1, self.adapter2]

    def test_fuzzy_prior_calculation(self) -> None:
        router = ExpertAdapterRouter(self.model, self.library, in_features=32)

        # 1. Text has statutes
        priors_statute = router.compute_fuzzy_priors("According to § 535 BGB...")
        # Index 0 is statute recall, index 1 is logic
        self.assertGreater(priors_statute[0].item(), priors_statute[1].item())

        # 2. Text has reasoning/FIRT
        priors_logic = router.compute_fuzzy_priors(
            "We analyze using FIRT reasoning traces"
        )
        self.assertGreater(priors_logic[1].item(), priors_logic[0].item())

    def test_hook_registration(self) -> None:
        router = ExpertAdapterRouter(self.model, self.library, in_features=32)
        self.assertEqual(len(router.hooks), 0)

        # Register hooks
        router.register_hooks()
        self.assertEqual(len(router.hooks), 1)

        # Unregister hooks
        router.unregister_hooks()
        self.assertEqual(len(router.hooks), 0)

    def test_forward_pass_with_routing(self) -> None:
        router = ExpertAdapterRouter(self.model, self.library, in_features=32)
        router.register_hooks()

        # Mock static active routing to only use adapter 1 (statute)
        priors = torch.tensor([1.0, 0.0])
        router.set_active_routing(priors)

        # Input tensor
        x = torch.ones(1, 4, 32)  # batch=1, seq_len=4, in_dim=32

        # 1. Standard forward pass through base model (without hooks)
        # To get the unadapted output, we can unregister hooks
        router.unregister_hooks()
        with torch.no_grad():
            output_base = self.model(x)

        # 2. Forward pass with active routing
        router.register_hooks()
        with torch.no_grad():
            output_adapted = self.model(x)

        # Check that adapter is applied
        self.assertFalse(torch.allclose(output_base, output_adapted))

        # Check mathematically:
        # For adapter 1: lora_B is 16x2 of 0.1, lora_A is 2x32 of 0.1
        # Input x is all ones of shape [1, 4, 32]
        # x_proj = x @ lora_A.t() -> shape [1, 4, 2].
        # Each entry of x_proj is sum_{k=1}^{32} 1.0 * 0.1 = 3.2
        # y_proj = x_proj @ lora_B.t() -> shape [1, 4, 16].
        # Each entry of y_proj is sum_{r=1}^2 3.2 * 0.1 = 0.64
        # Since weight prior is 1.0, adapted output should be output_base + 0.64
        expected_diff = torch.ones(1, 4, 16) * 0.64
        self.assertTrue(
            torch.allclose(output_adapted - output_base, expected_diff, rtol=1e-5)
        )

        router.unregister_hooks()


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