"""精度验证测试套件（Issue #39）。

覆盖：
- MAE/MRE/Max Error 误差计算工具
- PyTorch 参考实现对比
- 精度容差自动验证
- 数值稳定性测试
- 回归测试场景
"""

from __future__ import annotations

import numpy as np
import pytest
import torch

from sagellm_backend.dtype_conversion import (
    QuantizationMode,
    calculate_conversion_metrics,
    compute_fp8_quantization_params,
    dequantize_from_fp8_torch,
    quantize_to_fp8_torch,
    validate_conversion,
)


class TestPrecisionMetrics:
    """误差计算工具测试。"""

    def test_calculate_conversion_metrics_torch(self) -> None:
        """验证 torch 张量的 MAE/MRE/Max Error 计算。"""
        original = torch.tensor([1.0, -2.0, 4.0], dtype=torch.float32)
        reconstructed = torch.tensor([1.1, -1.8, 3.7], dtype=torch.float32)

        metrics = calculate_conversion_metrics(original, reconstructed, "torch_case")

        expected_mae = torch.mean(torch.abs(original - reconstructed)).item()
        expected_mre = torch.mean(torch.abs((original - reconstructed) / (original + 1e-8))).item()
        expected_max_error = torch.max(torch.abs(original - reconstructed)).item()

        assert metrics.conversion_path == "torch_case"
        assert metrics.mae == pytest.approx(expected_mae, rel=1e-6)
        assert metrics.mre == pytest.approx(expected_mre, rel=1e-6)
        assert metrics.max_error == pytest.approx(expected_max_error, rel=1e-6)

    def test_calculate_conversion_metrics_numpy(self) -> None:
        """验证 numpy 数组的 MAE/MRE/Max Error 计算。"""
        original = np.array([1.0, -2.0, 4.0], dtype=np.float32)
        reconstructed = np.array([1.1, -1.8, 3.7], dtype=np.float32)

        metrics = calculate_conversion_metrics(original, reconstructed, "numpy_case")

        expected_mae = float(np.mean(np.abs(original - reconstructed)))
        expected_mre = float(np.mean(np.abs((original - reconstructed) / (original + 1e-8))))
        expected_max_error = float(np.max(np.abs(original - reconstructed)))

        assert metrics.conversion_path == "numpy_case"
        assert metrics.mae == pytest.approx(expected_mae, rel=1e-6)
        assert metrics.mre == pytest.approx(expected_mre, rel=1e-6)
        assert metrics.max_error == pytest.approx(expected_max_error, rel=1e-6)


class TestToleranceValidation:
    """精度容差自动验证测试。"""

    def test_validate_conversion_passes_within_threshold(self) -> None:
        """误差在阈值内时通过验证。"""
        base = torch.linspace(-1.0, 1.0, steps=256, dtype=torch.float32)
        noise = torch.randn_like(base) * 1e-4
        reconstructed = base + noise

        metrics = validate_conversion(
            base,
            reconstructed,
            "tolerance_ok",
            mae_threshold=0.01,
            mre_threshold=0.05,
        )

        assert metrics.mae < 0.01
        assert metrics.mre < 0.05

    def test_validate_conversion_fails_mre(self) -> None:
        """MRE 超阈值时应抛错。"""
        original = torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float32)
        reconstructed = torch.tensor([1.0, 1.0, 1.0, 1.5], dtype=torch.float32)

        with pytest.raises(ValueError, match="MRE .* exceeds threshold"):
            validate_conversion(
                original,
                reconstructed,
                "tolerance_fail_mre",
                mae_threshold=1.0,
                mre_threshold=0.05,
            )


class TestNumericalStability:
    """数值稳定性测试。"""

    def test_metrics_near_zero_stable(self) -> None:
        """接近零值输入不会产生 NaN/Inf。"""
        original = torch.tensor([0.0, 1e-9, -1e-9, 1e-7], dtype=torch.float32)
        reconstructed = torch.tensor([1e-10, 2e-9, -2e-9, 2e-7], dtype=torch.float32)

        metrics = calculate_conversion_metrics(original, reconstructed, "near_zero")

        assert np.isfinite(metrics.mae)
        assert np.isfinite(metrics.mre)
        assert np.isfinite(metrics.max_error)

    def test_fp8_quantization_stability_large_dynamic_range(self) -> None:
        """大动态范围下量化路径可稳定执行并保持合理误差。"""
        torch.manual_seed(39)
        exponents = torch.randint(low=-4, high=4, size=(2048,), dtype=torch.int32)
        signs = torch.where(torch.rand(2048) > 0.5, 1.0, -1.0)
        original = signs * torch.pow(torch.tensor(10.0), exponents.float())

        params = compute_fp8_quantization_params(original, mode=QuantizationMode.BLOCK_WISE)
        quantized = quantize_to_fp8_torch(original, params)
        reconstructed = dequantize_from_fp8_torch(quantized, params)
        metrics = calculate_conversion_metrics(original, reconstructed, "fp8_dynamic_range")

        assert torch.isfinite(reconstructed).all()
        assert metrics.mae < 1.2
        assert metrics.max_error < 6.0


class TestRegressionSuite:
    """回归测试套件。"""

    @pytest.mark.parametrize(
        ("seed", "shape", "mode", "mae_threshold", "mre_threshold"),
        [
            (39, (64, 64), QuantizationMode.PER_TENSOR, 0.15, 0.0125),
            (40, (128, 32), QuantizationMode.BLOCK_WISE, 0.15, 0.0125),
            (41, (1024,), QuantizationMode.BLOCK_WISE, 0.15, 0.0125),
        ],
    )
    def test_fp8_reference_regression(
        self,
        seed: int,
        shape: tuple[int, ...],
        mode: QuantizationMode,
        mae_threshold: float,
        mre_threshold: float,
    ) -> None:
        """固定种子回归：对比 PyTorch 参考算子路径。"""
        torch.manual_seed(seed)
        # 使用远离 0 的有界分布，避免相对误差被近零分母放大
        magnitude = torch.rand(*shape, dtype=torch.float32) * 4 + 1
        sign = torch.where(torch.rand(*shape) > 0.5, 1.0, -1.0)
        original = magnitude * sign

        params = compute_fp8_quantization_params(original, mode=mode)
        quantized = quantize_to_fp8_torch(original, params)
        reconstructed = dequantize_from_fp8_torch(quantized, params)

        metrics = validate_conversion(
            original,
            reconstructed,
            f"FP32->FP8({mode.value})->FP32",
            mae_threshold=mae_threshold,
            mre_threshold=mre_threshold,
        )

        # 与“参考实现”一致：误差直接由 PyTorch 基础算子计算
        ref_mae = torch.mean(torch.abs(original - reconstructed)).item()
        ref_mre = torch.mean(torch.abs((original - reconstructed) / (original + 1e-8))).item()

        assert metrics.mae == pytest.approx(ref_mae, rel=1e-6)
        assert metrics.mre == pytest.approx(ref_mre, rel=1e-6)
