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 (
    ParameterDiagnostics,
    ParameterDiagnosticsConfig,
    TrainingSession,
    validate_record,
)
from tests.service_test_utils import unique_project_name


@unittest.skipIf(torch is None or nn is None, "PyTorch is required for parameter diagnostics tests.")
class ParameterDiagnosticsTests(unittest.TestCase):
    def test_parameter_diagnostics_writes_summaries_and_distributions(self) -> None:
        class SmallNet(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.encoder = nn.Linear(4, 5)
                self.head = nn.Linear(5, 2)

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

        torch.manual_seed(13)
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("parameter-diagnostics"),
                storage_dir=storage_dir,
            )
            model = SmallNet()
            optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
            diagnostics = ParameterDiagnostics(
                session,
                model,
                optimizer,
                config=ParameterDiagnosticsConfig(
                    batch_size=16,
                    capture_distributions=True,
                    distribution_types=("value", "gradient", "update"),
                    histogram_bins=8,
                ),
            )

            loss_fn = nn.CrossEntropyLoss()
            for step in range(2):
                inputs = torch.randn(6, 4)
                targets = torch.tensor([0, 1, 0, 1, 1, 0], dtype=torch.long)
                optimizer.zero_grad(set_to_none=True)
                loss = loss_fn(model(inputs), targets)
                loss.backward()
                optimizer.step()
                diagnostics.snapshot(step, epoch=0, metadata={"test": "parameter-diagnostics"})

            summaries = diagnostics.read_summaries()
            distributions = diagnostics.read_distributions()
            latest_weight_summaries = diagnostics.read_summaries(
                parameter_role="weight",
                latest_only=True,
            )
            diagnostics.close()

            self.assertTrue(summaries)
            self.assertTrue(distributions)
            self.assertTrue(latest_weight_summaries)
            self.assertEqual({summary.step for summary in latest_weight_summaries}, {1})
            self.assertIn("weight", {summary.parameter_role for summary in summaries})
            self.assertIn("bias", {summary.parameter_role for summary in summaries})
            self.assertTrue(any(summary.update_norm is not None for summary in summaries if summary.step == 1))
            self.assertTrue(any(item.distribution_type == "value" for item in distributions))
            self.assertTrue(any(item.distribution_type == "gradient" for item in distributions))
            self.assertTrue(any(item.distribution_type == "update" for item in distributions))
            self.assertTrue(all(len(item.bin_edges) == len(item.counts) + 1 for item in distributions))
            validate_record(summaries[0])
            validate_record(distributions[0])


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