"""
Tests for Dynamic Precision Allocation.

Verifies that:
1. Quantization/dequantization is lossless for FP16
2. Quantization error is bounded for lower precisions
3. The precision planner respects memory budgets
4. Block-wise quantization correctly packs/unpacks values
"""

import pytest
import torch

from dokodemo_ai.quantization.dynamic import DynamicQuantizer, Precision, PrecisionPlan


class TestPrecision:
    def test_bytes_per_element(self):
        assert Precision.FP16.bytes_per_element == 2.0
        assert Precision.INT8.bytes_per_element == 1.0
        assert Precision.INT4.bytes_per_element == 0.5
        assert Precision.INT2.bytes_per_element == 0.25

    def test_compression_ratio(self):
        assert Precision.FP16.compression_ratio == 1.0
        assert Precision.INT8.compression_ratio == 2.0
        assert Precision.INT4.compression_ratio == 4.0
        assert Precision.INT2.compression_ratio == 8.0


class TestQuantizationFP16:
    """FP16 quantization should be effectively lossless."""

    def test_fp16_roundtrip(self):
        q = DynamicQuantizer()
        x = torch.randn(256, 256)
        quantized, meta = q.quantize(x, Precision.FP16)
        recovered = q.dequantize(quantized, meta)
        assert torch.allclose(x.half(), recovered, atol=1e-3)

    def test_fp16_meta_tag(self):
        q = DynamicQuantizer()
        x = torch.randn(64)
        _, meta = q.quantize(x, Precision.FP16)
        assert meta["precision"] == "fp16"


class TestQuantizationINT8:
    """INT8 should be low-error and halve the size."""

    def test_int8_roundtrip_shape(self):
        q = DynamicQuantizer(block_size=32)
        x = torch.randn(128, 64)
        quantized, meta = q.quantize(x, Precision.INT8)
        recovered = q.dequantize(quantized, meta)
        assert recovered.shape == x.shape

    def test_int8_error_bounded(self):
        q = DynamicQuantizer(block_size=32)
        x = torch.randn(512)
        quantized, meta = q.quantize(x, Precision.INT8)
        recovered = q.dequantize(quantized, meta)
        # INT8 should be < 1% relative error for typical activations
        max_err = (recovered.float() - x.float()).abs().max().item()
        x_range = x.float().abs().max().item()
        relative_err = max_err / (x_range + 1e-8)
        assert relative_err < 0.05, f"Relative error too high: {relative_err:.4f}"

    def test_int8_meta_tag(self):
        q = DynamicQuantizer(block_size=64)
        x = torch.ones(128)
        _, meta = q.quantize(x, Precision.INT8)
        assert meta["precision"] == "INT8"
        assert "scale" in meta
        assert "zero_point" in meta


class TestQuantizationINT4:
    """INT4 should recover shape and be within reasonable error."""

    def test_int4_roundtrip_shape(self):
        q = DynamicQuantizer(block_size=64)
        x = torch.randn(256)
        quantized, meta = q.quantize(x, Precision.INT4)
        recovered = q.dequantize(quantized, meta)
        assert recovered.shape == x.shape

    def test_int4_reduces_storage(self):
        q = DynamicQuantizer(block_size=64)
        x = torch.randn(512)
        quantized, meta = q.quantize(x, Precision.INT4)
        # Packed INT4 should take ~1/4 the elements (2 values per byte)
        # Meta tensors (scale, zero_point) are small
        original_bytes = x.numel() * 2  # FP16
        packed_bytes = quantized.numel()  # uint8, each holds 2 int4 values
        assert packed_bytes <= original_bytes // 2 + 1  # Allow for padding

    def test_int4_meta_tag(self):
        q = DynamicQuantizer(block_size=32)
        x = torch.randn(64)
        _, meta = q.quantize(x, Precision.INT4)
        assert meta["precision"] == "INT4"

    def test_int4_constant_tensor(self):
        """Constant tensor should quantize/dequantize near-perfectly."""
        q = DynamicQuantizer(block_size=32)
        x = torch.ones(256) * 3.14
        quantized, meta = q.quantize(x, Precision.INT4)
        recovered = q.dequantize(quantized, meta)
        # For a constant tensor, all blocks should be identical
        diff = (recovered.float() - x.float()).abs().mean().item()
        assert diff < 0.1, f"Constant tensor dequantization error: {diff}"


class TestQuantizationINT2:
    """INT2 is aggressive but should still work structurally."""

    def test_int2_roundtrip_shape(self):
        q = DynamicQuantizer(block_size=64)
        x = torch.randn(256)
        quantized, meta = q.quantize(x, Precision.INT2)
        recovered = q.dequantize(quantized, meta)
        assert recovered.shape == x.shape

    def test_int2_meta_tag(self):
        q = DynamicQuantizer(block_size=64)
        x = torch.randn(64)
        _, meta = q.quantize(x, Precision.INT2)
        assert meta["precision"] == "INT2"


class TestPrecisionPlanner:
    """Tests for the memory-budget precision planner."""

    def test_plan_respects_budget(self):
        """Total compressed size should not exceed budget."""
        q = DynamicQuantizer()
        # 10 layers, each 1GB FP16 original size
        layer_sizes = [1 * 1024**3] * 10
        budget = 4 * 1024**3  # 4 GB

        plan = q.create_plan(layer_sizes, budget)
        assert plan.total_size_bytes <= budget, (
            f"Plan ({plan.total_size_bytes / 1e9:.2f} GB) exceeds "
            f"budget ({budget / 1e9:.2f} GB)"
        )

    def test_plan_uses_higher_precision_for_sensitive(self):
        """More sensitive layers should get higher precision."""
        q = DynamicQuantizer()
        n = 10
        layer_sizes = [512 * 1024**2] * n  # 512 MB each
        budget = 6 * 1024**3  # 6 GB — enough for some FP16

        # Force first/last layers to be most sensitive
        sensitivities = [1.0] + [0.2] * (n - 2) + [1.0]
        plan = q.create_plan(layer_sizes, budget, sensitivities)

        first_prec = plan.layer_configs[0].precision.value
        last_prec = plan.layer_configs[-1].precision.value
        mid_prec = plan.layer_configs[n // 2].precision.value

        # First/last should have >= precision as middle
        assert first_prec >= mid_prec, (
            f"First layer ({first_prec}-bit) should be >= middle ({mid_prec}-bit)"
        )
        assert last_prec >= mid_prec

    def test_plan_all_layers_assigned(self):
        """Every layer must have a precision assignment."""
        q = DynamicQuantizer()
        n = 20
        layer_sizes = [200 * 1024**2] * n
        budget = 2 * 1024**3
        plan = q.create_plan(layer_sizes, budget)
        assert len(plan.layer_configs) == n

    def test_heuristic_sensitivities_u_shaped(self):
        """Heuristic sensitivities should be highest at edges."""
        q = DynamicQuantizer()
        n = 20
        sens = q._heuristic_sensitivities(n)

        assert sens[0] > sens[n // 2], "First layer should score higher than middle"
        assert sens[-1] > sens[n // 2], "Last layer should score higher than middle"
        assert all(0.0 <= s <= 1.0 for s in sens), "Scores must be in [0, 1]"

    def test_tight_budget_forces_low_precision(self):
        """Very tight budget should force INT2 on all layers."""
        q = DynamicQuantizer()
        layer_sizes = [1 * 1024**3] * 5  # 5 x 1 GB = 5 GB FP16
        budget = 700 * 1024**2  # 700 MB — very tight

        plan = q.create_plan(layer_sizes, budget)
        # All layers should be INT2 (lowest precision)
        for cfg in plan.layer_configs:
            assert cfg.precision == Precision.INT2, (
                f"Expected INT2, got {cfg.precision}"
            )

    def test_ample_budget_allows_fp16(self):
        """If budget is sufficient, at least some layers should be FP16."""
        q = DynamicQuantizer()
        layer_sizes = [100 * 1024**2] * 4  # 4 x 100 MB = 400 MB FP16
        budget = 2 * 1024**3  # 2 GB — very ample

        plan = q.create_plan(layer_sizes, budget)
        precisions = [c.precision for c in plan.layer_configs]
        assert Precision.FP16 in precisions, (
            "Ample budget should allow FP16 for some layers"
        )

    def test_plan_summary_runs(self):
        """Plan summary should not raise exceptions."""
        q = DynamicQuantizer()
        plan = q.create_plan([100 * 1024**2] * 4, 1 * 1024**3)
        summary = plan.summary()
        assert "Precision plan" in summary
        assert "GB" in summary
