from __future__ import annotations

import sys
import tempfile
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"

if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

try:
    import torch
    import torch.nn as nn
except Exception:  # pragma: no cover - exercised by skip condition.
    torch = None
    nn = None

from quark import HookEngine, HookEngineConfig, HookScheduleConfig, LayerFilterConfig, TrainingSession
from tests.service_test_utils import unique_project_name


@unittest.skipIf(torch is None or nn is None, "PyTorch is required for hook engine tests.")
class HookEngineTests(unittest.TestCase):
    def test_hook_engine_emits_activation_and_gradient_summaries_for_common_layers(self) -> None:
        class SmallConvNet(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.relu = nn.ReLU()
                self.flatten = nn.Flatten()
                self.norm = nn.LayerNorm(4 * 8 * 8)
                self.head = nn.Linear(4 * 8 * 8, 3)

            def forward(self, inputs: torch.Tensor) -> torch.Tensor:
                outputs = self.conv(inputs)
                outputs = self.relu(outputs)
                outputs = self.flatten(outputs)
                outputs = self.norm(outputs)
                return self.head(outputs)

        torch.manual_seed(7)
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("hooks-demo"),
                storage_dir=storage_dir,
            )
            model = SmallConvNet()
            optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
            loss_fn = nn.CrossEntropyLoss()
            config = HookEngineConfig(
                batch_size=8,
                filters=LayerFilterConfig(
                    include_types=("Conv2d", "ReLU", "LayerNorm", "Linear"),
                ),
            )
            hooks = HookEngine(session, model, optimizer, config=config)

            initial_weight = model.head.weight.detach().clone()

            for step in range(2):
                inputs = torch.randn(4, 3, 8, 8)
                targets = torch.tensor([0, 1, 2, 1], dtype=torch.long)
                optimizer.zero_grad(set_to_none=True)
                with hooks.step(step, epoch=0):
                    logits = model(inputs)
                    loss = loss_fn(logits, targets)
                    loss.backward()
                optimizer.step()

            summaries = hooks.read()
            hooks.close()

            self.assertTrue(torch.any(model.head.weight.detach().ne(initial_weight)).item())

            summary_types = {summary.summary_type for summary in summaries}
            module_names = {summary.module_name for summary in summaries}

            self.assertEqual(summary_types, {"activation", "gradient"})
            self.assertTrue({"conv", "relu", "norm", "head"}.issubset(module_names))
            self.assertEqual({summary.step for summary in summaries}, {0, 1})
            self.assertTrue(all(summary.norm is not None and summary.norm >= 0.0 for summary in summaries))

            relu_activations = [
                summary for summary in summaries if summary.module_name == "relu" and summary.summary_type == "activation"
            ]
            self.assertTrue(relu_activations)
            self.assertTrue(all(summary.dead_unit_ratio is not None for summary in relu_activations))

    def test_hook_engine_respects_schedule_and_layer_filters(self) -> None:
        class FilteredNet(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.encoder = nn.Linear(4, 8)
                self.activation = nn.Tanh()
                self.head = nn.Linear(8, 2)

            def forward(self, inputs: torch.Tensor) -> torch.Tensor:
                outputs = self.encoder(inputs)
                outputs = self.activation(outputs)
                return self.head(outputs)

        torch.manual_seed(11)
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("filtered-hooks"),
                storage_dir=storage_dir,
            )
            model = FilteredNet()
            optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
            loss_fn = nn.MSELoss()
            config = HookEngineConfig(
                batch_size=4,
                schedule=HookScheduleConfig(every_n_steps=2, start_step=1),
                filters=LayerFilterConfig(
                    include_types=("Linear", "Tanh"),
                    exclude_names=("head",),
                ),
            )
            hooks = HookEngine(session, model, optimizer, config=config)

            for step in range(4):
                inputs = torch.randn(3, 4)
                targets = torch.randn(3, 2)
                optimizer.zero_grad(set_to_none=True)
                with hooks.step(step, epoch=0):
                    outputs = model(inputs)
                    loss = loss_fn(outputs, targets)
                    loss.backward()
                optimizer.step()

            summaries = hooks.read()
            hooks.close()

            self.assertEqual({summary.step for summary in summaries}, {1, 3})
            self.assertEqual({summary.module_name for summary in summaries}, {"encoder", "activation"})
            self.assertTrue(all(summary.module_name != "head" for summary in summaries))

            tanh_activations = [
                summary
                for summary in summaries
                if summary.module_name == "activation" and summary.summary_type == "activation"
            ]
            self.assertTrue(tanh_activations)
            self.assertTrue(
                all(summary.saturation_ratio is not None for summary in tanh_activations)
            )


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