"""Integration of quantization analysis with CAAE dashboard"""

from fastapi import APIRouter
from .quantization_analysis import QuantizationAnalyzer
from .schemas import QuantizationRequest, QuantizationResponse

router = APIRouter(prefix="/quantization", tags=["quantization"])


@router.post("/analyze", response_model=QuantizationResponse)
async def analyze_quantization(request: QuantizationRequest):
    """Analyze quantization options for given model configuration"""
    analyzer = QuantizationAnalyzer()

    # Run analysis
    metrics = analyzer.analyze_model_compression(
        request.model_size_gb, request.target_dtype
    )

    recommendation = analyzer.recommend_quantization(
        request.model_size_gb, request.memory_budget_gb
    )

    kv_compression = analyzer.calculate_kv_cache_compression(
        request.batch_size, request.seq_len, request.hidden_size, request.num_layers
    )

    return QuantizationResponse(
        compression_ratio=metrics.compression_ratio,
        memory_savings_mb=metrics.memory_savings_mb,
        latency_impact_ms=metrics.latency_impact_ms,
        accuracy_degradation=metrics.accuracy_degradation,
        recommendation=recommendation,
        kv_cache_compression=kv_compression,
    )


@router.get("/recommendations/{model_name}")
async def get_quantization_recommendations(model_name: str):
    """Get quantization recommendations for common models"""
    model_configs = {
        "mistral-7b": {"size_gb": 14, "memory_budget_gb": 24},
        "llama2-13b": {"size_gb": 26, "memory_budget_gb": 40},
        "codellama-34b": {"size_gb": 68, "memory_budget_gb": 80},
    }

    if model_name not in model_configs:
        return {"error": "Model not found"}

    config = model_configs[model_name]
    analyzer = QuantizationAnalyzer()

    return {
        "model": model_name,
        "recommendation": analyzer.recommend_quantization(
            config["size_gb"], config["memory_budget_gb"]
        ),
        "analysis": analyzer.analyze_model_compression(
            config["size_gb"],
            analyzer.recommend_quantization(
                config["size_gb"], config["memory_budget_gb"]
            ),
        ).__dict__,
    }
