"""Calibration module for A2RAG package."""
from __future__ import annotations
from typing import List, Dict, Optional
from .models import CalibrationResult


def run_calibration(
    data: List[Dict],
    api_client,
    domain: Optional[str] = None,
) -> CalibrationResult:
    """
    Find optimal thresholds via API calibration endpoint.
    Falls back to local grid search if API unavailable.
    """
    try:
        result = api_client.post("/calibrate", {
            "data": data,
            "domain": domain,
        })
        return CalibrationResult(
            domain=domain or "generic",
            tau_evidence=result["tau_evidence"],
            tau_completeness=result["tau_completeness"],
            tau_abstain=result["tau_abstain"],
            expected_accuracy=result["accuracy"],
            expected_uar=result["UAR"],
            expected_ors=result.get("ORS", 0),
            dataset_size=len(data),
            recommendation=result.get("recommendation", ""),
        )
    except Exception:
        return _local_calibration(data, domain)


def _local_calibration(data: List[Dict], domain: Optional[str]) -> CalibrationResult:
    """Local grid search calibration."""
    from itertools import product as iproduct
    import sys, os
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

    try:
        from core.decision_layer import A2RAGDecisionLayer
    except ImportError:
        return CalibrationResult(
            domain=domain or "generic",
            tau_evidence=0.40, tau_completeness=0.80, tau_abstain=0.20,
            expected_accuracy=0, expected_uar=0, expected_ors=0,
            dataset_size=len(data),
            recommendation="Install a2rag server for calibration",
        )

    best = None
    best_score = -1

    for tau_e, tau_c, tau_a in iproduct(
        [0.30, 0.40, 0.55, 0.70],
        [0.65, 0.75, 0.80, 0.90],
        [0.15, 0.20, 0.25],
    ):
        layer = A2RAGDecisionLayer(
            tau_evidence=tau_e, tau_completeness=tau_c, tau_abstain=tau_a
        )
        results = []
        for item in data:
            r = layer.decide(item["query"], item["contexts"], item["draft_answer"])
            results.append({
                "expected": item["label"],
                "predicted": r.action,
                "correct": r.action == item["label"],
            })
        n = len(results)
        unsafe = sum(1 for r in results
                     if r["predicted"] == "answer" and r["expected"] != "answer")
        correct = sum(1 for r in results if r["correct"])
        uar = unsafe / n
        acc = correct / n
        score = acc * (1 - uar)
        if uar == 0 and score > best_score:
            best_score = score
            best = {"tau_e": tau_e, "tau_c": tau_c, "tau_a": tau_a,
                    "acc": acc, "uar": uar}

    if best is None:
        best = {"tau_e": 0.40, "tau_c": 0.80, "tau_a": 0.20, "acc": 0, "uar": 0}

    return CalibrationResult(
        domain=domain or "generic",
        tau_evidence=best["tau_e"],
        tau_completeness=best["tau_c"],
        tau_abstain=best["tau_a"],
        expected_accuracy=best["acc"],
        expected_uar=best["uar"],
        expected_ors=best["acc"] * (1 - best["uar"]),
        dataset_size=len(data),
        recommendation=f"Use tau_evidence={best['tau_e']}, tau_completeness={best['tau_c']}",
    )
