"""Tests for quantization functionality (Task3.1)."""

from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
import pytest

from sagellm_compression.quantization import (
    Calibrator,
    Exporter,
    FP8Config,
    Quantizer,
    W4A16Config,
    W8A8Config,
)
from sagellm_compression.quantization.configs import CalibrationMethod, QuantizationType


class TestQuantizationConfigs:
    """Test quantization configuration classes."""

    def test_w8a8_config_default(self) -> None:
        """Test W8A8 config with default parameters."""
        config = W8A8Config()
        assert config.quantization_type == QuantizationType.W8A8
        assert config.per_channel is True
        assert config.symmetric is True
        assert config.quantize_weights is True
        assert config.quantize_activations is True

    def test_w4a16_config_default(self) -> None:
        """Test W4A16 config with default parameters."""
        config = W4A16Config()
        assert config.quantization_type == QuantizationType.W4A16
        assert config.per_channel is True
        assert config.symmetric is False  # W4A16 uses asymmetric
        assert config.quantize_weights is True
        assert config.quantize_activations is False  # Activations stay in FP16
        assert config.group_size == 128

    def test_fp8_config_default(self) -> None:
        """Test FP8 config with default parameters."""
        config = FP8Config()
        assert config.quantization_type == QuantizationType.FP8
        assert config.per_channel is True
        assert config.symmetric is True
        assert config.fp8_format == "e4m3"

    def test_custom_calibration_method(self) -> None:
        """Test custom calibration method."""
        config = W8A8Config(calibration_method=CalibrationMethod.PERCENTILE)
        assert config.calibration_method == CalibrationMethod.PERCENTILE

    def test_custom_num_samples(self) -> None:
        """Test custom number of calibration samples."""
        config = W8A8Config(num_calibration_samples=256)
        assert config.num_calibration_samples == 256

    def test_w4a16_custom_group_size(self) -> None:
        """Test W4A16 with custom group size."""
        config = W4A16Config(group_size=64)
        assert config.group_size == 64


class TestCalibrator:
    """Test Calibrator class."""

    def test_calibrator_initialization(self) -> None:
        """Test calibrator initialization."""
        config = W8A8Config()
        calibrator = Calibrator(config)
        assert calibrator.config == config
        assert len(calibrator.samples) == 0
        assert calibrator.params is None

    def test_collect_samples(self) -> None:
        """Test collecting calibration samples."""
        config = W8A8Config(num_calibration_samples=10)
        calibrator = Calibrator(config)

        # Collect some samples
        for i in range(5):
            sample = np.random.randn(32, 64).astype(np.float32)
            calibrator.collect(sample)

        assert len(calibrator.samples) == 5

    def test_compute_params_minmax(self) -> None:
        """Test computing quantization parameters with minmax calibration."""
        config = W8A8Config(calibration_method=CalibrationMethod.MINMAX)
        calibrator = Calibrator(config)

        # Collect samples with known range
        samples = [
            np.array([[-2.0, -1.0], [0.0, 1.0]], dtype=np.float32),
            np.array([[-1.5, 0.5], [1.5, 2.0]], dtype=np.float32),
        ]
        for sample in samples:
            calibrator.collect(sample)

        params = calibrator.compute_params(bits=8)
        assert params is not None
        assert params.min_val == pytest.approx(-2.0, abs=1e-5)
        assert params.max_val == pytest.approx(2.0, abs=1e-5)

    def test_compute_params_symmetric(self) -> None:
        """Test symmetric quantization (zero_point = 0)."""
        config = W8A8Config(symmetric=True)
        calibrator = Calibrator(config)

        sample = np.array([[-1.0, 0.0, 1.0]], dtype=np.float32)
        calibrator.collect(sample)

        params = calibrator.compute_params(bits=8)
        assert params.zero_point == 0  # Symmetric

    def test_compute_params_no_samples_raises(self) -> None:
        """Test that computing params without samples raises error."""
        config = W8A8Config()
        calibrator = Calibrator(config)

        with pytest.raises(ValueError, match="No samples collected"):
            calibrator.compute_params()

    def test_calibrator_reset(self) -> None:
        """Test resetting calibrator state."""
        config = W8A8Config()
        calibrator = Calibrator(config)

        sample = np.random.randn(10, 10).astype(np.float32)
        calibrator.collect(sample)
        calibrator.compute_params()

        calibrator.reset()
        assert len(calibrator.samples) == 0
        assert calibrator.params is None


class TestQuantizer:
    """Test Quantizer class."""

    def test_quantizer_initialization(self) -> None:
        """Test quantizer initialization."""
        config = W8A8Config()
        quantizer = Quantizer(config)
        assert quantizer.config == config

    def test_quantize_tensor_w8a8(self) -> None:
        """Test quantizing a single tensor with W8A8."""
        config = W8A8Config()
        quantizer = Quantizer(config)

        # Create a sample tensor
        tensor = np.array([[-1.0, 0.0, 1.0], [2.0, 3.0, 4.0]], dtype=np.float32)

        # Compute quantization params
        calibrator = Calibrator(config)
        calibrator.collect(tensor)
        params = calibrator.compute_params(bits=8)

        # Quantize
        quantized = quantizer.quantize_tensor(tensor, params, bits=8)
        assert quantized.data.dtype == np.int8
        assert quantized.original_shape == tensor.shape

        # Dequantize and check approximate reconstruction
        dequantized = quantized.dequantize()
        assert dequantized.shape == tensor.shape
        # Allow some quantization error
        assert np.allclose(dequantized, tensor, atol=0.1)

    def test_quantize_model_weights(self) -> None:
        """Test quantizing model weights."""
        config = W8A8Config()
        quantizer = Quantizer(config)

        # Simulate model weights
        model_weights = {
            "layer1.weight": np.random.randn(64, 128).astype(np.float32),
            "layer2.weight": np.random.randn(128, 256).astype(np.float32),
        }

        quantized_model = quantizer.quantize(model_weights)
        assert len(quantized_model.weights) == 2
        assert "layer1.weight" in quantized_model.weights
        assert "layer2.weight" in quantized_model.weights
        assert quantized_model.metadata["num_layers"] == 2

    def test_quantize_w4a16(self) -> None:
        """Test W4A16 quantization."""
        config = W4A16Config(group_size=64)
        quantizer = Quantizer(config)

        model_weights = {
            "weight": np.random.randn(128, 128).astype(np.float32),
        }

        quantized_model = quantizer.quantize(model_weights)
        assert len(quantized_model.weights) == 1
        # W4A16 uses 4-bit weights
        weight = quantized_model.weights["weight"]
        assert weight.data.dtype == np.int8  # Stored as int8


class TestExporter:
    """Test Exporter class."""

    def test_exporter_initialization(self) -> None:
        """Test exporter initialization."""
        exporter = Exporter()
        assert exporter is not None

    def test_export_and_load_pickle(self) -> None:
        """Test exporting and loading in pickle format."""
        config = W8A8Config()
        quantizer = Quantizer(config)

        model_weights = {
            "layer.weight": np.random.randn(32, 64).astype(np.float32),
        }
        quantized_model = quantizer.quantize(model_weights)

        exporter = Exporter()

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "model.pkl"

            # Export
            exporter.export(quantized_model, output_path, format="pickle")
            assert output_path.exists()

            # Load
            loaded_model = exporter.load(output_path, format="pickle")
            assert len(loaded_model.weights) == len(quantized_model.weights)
            assert loaded_model.config.quantization_type == config.quantization_type

    def test_export_metadata(self) -> None:
        """Test exporting model metadata to JSON."""
        config = W8A8Config()
        quantizer = Quantizer(config)

        model_weights = {
            "layer.weight": np.random.randn(16, 32).astype(np.float32),
        }
        quantized_model = quantizer.quantize(model_weights)

        exporter = Exporter()

        with tempfile.TemporaryDirectory() as tmpdir:
            metadata_path = Path(tmpdir) / "metadata.json"

            exporter.export_metadata(quantized_model, metadata_path)
            assert metadata_path.exists()

            # Check content
            import json

            with open(metadata_path) as f:
                metadata = json.load(f)

            assert metadata["quantization_type"] == "w8a8"
            assert metadata["num_layers"] == 1
            assert "layer.weight" in metadata["layer_names"]

    def test_load_nonexistent_file_raises(self) -> None:
        """Test loading non-existent file raises error."""
        exporter = Exporter()

        with pytest.raises(FileNotFoundError):
            exporter.load("/nonexistent/path/model.pkl")


class TestEndToEnd:
    """End-to-end integration tests."""

    def test_full_quantization_pipeline_w8a8(self) -> None:
        """Test full quantization pipeline with W8A8."""
        # 1. Create config
        config = W8A8Config()

        # 2. Quantize model
        quantizer = Quantizer(config)
        model_weights = {
            "fc1.weight": np.random.randn(128, 256).astype(np.float32),
            "fc2.weight": np.random.randn(256, 512).astype(np.float32),
        }
        quantized_model = quantizer.quantize(model_weights)

        # 3. Export
        exporter = Exporter()
        with tempfile.TemporaryDirectory() as tmpdir:
            model_path = Path(tmpdir) / "model_int8.pkl"
            exporter.export(quantized_model, model_path)

            # 4. Load
            loaded_model = exporter.load(model_path)

            # 5. Verify
            assert len(loaded_model.weights) == 2
            assert set(loaded_model.weights.keys()) == {"fc1.weight", "fc2.weight"}

            # Dequantize and check reconstruction quality
            for layer_name, original_weight in model_weights.items():
                quantized_weight = loaded_model.weights[layer_name]
                dequantized = quantized_weight.dequantize()
                # Check shape matches
                assert dequantized.shape == original_weight.shape
                # Check quantization introduces bounded error
                error = np.abs(dequantized - original_weight).mean()
                assert error < 0.5  # Reasonable quantization error for INT8

    def test_full_quantization_pipeline_w4a16(self) -> None:
        """Test full quantization pipeline with W4A16."""
        config = W4A16Config(group_size=64)
        quantizer = Quantizer(config)

        model_weights = {
            "weight": np.random.randn(64, 64).astype(np.float32),
        }
        quantized_model = quantizer.quantize(model_weights)

        exporter = Exporter()
        with tempfile.TemporaryDirectory() as tmpdir:
            model_path = Path(tmpdir) / "model_int4.pkl"
            exporter.export(quantized_model, model_path)
            loaded_model = exporter.load(model_path)

            assert len(loaded_model.weights) == 1
            assert loaded_model.config.quantization_type == QuantizationType.W4A16


def test_quantization_placeholder() -> None:
    """Placeholder test to ensure test suite runs."""
    assert True, "Quantization module placeholder test"
