"""
Chuẩn metrics khi lưu / POST lên DB (bảng mlflow.metrics).
Bắt buộc: Accuracy, Precision, Recall, F1-Score, Confusion Matrix.
"""

from typing import Dict, Any, List, Optional, Union
import numpy as np

# Các metric chuẩn bắt buộc khi POST /api/v1/metrics
STANDARD_METRIC_KEYS = [
    "accuracy",
    "precision",
    "recall",
    "f1_score",
    "confusion_matrix",
]

# Optional metrics (backend hỗ trợ)
OPTIONAL_METRIC_KEYS = [
    "auc_macro",
    "auc_micro",
    "ap_macro",
    "ap_micro",
    "roc_curve",
    "pr_curve",
]


def _ensure_confusion_matrix_list(cm: Any) -> Optional[List[List[int]]]:
    """Chuyển confusion matrix sang list of list (JSON-serializable)."""
    if cm is None:
        return None
    if isinstance(cm, np.ndarray):
        return cm.astype(int).tolist()
    if isinstance(cm, list):
        if cm and isinstance(cm[0], (list, np.ndarray)):
            return [[int(x) for x in row] for row in cm]
        return cm
    return None


def build_metrics_payload(
    model_id: int,
    dataset_split_id: int,
    experiment_run_id: int,
    metrics: Dict[str, Any],
    *,
    require_standard: bool = True,
) -> Dict[str, Any]:
    """
    Tạo payload chuẩn để POST /api/v1/metrics.

    Chuẩn bắt buộc: accuracy, precision, recall, f1_score, confusion_matrix.
    Các giá trị phải trong [0, 1] (trừ confusion_matrix là 2D list int).

    Args:
        model_id: ID model
        dataset_split_id: ID dataset split
        experiment_run_id: ID experiment run
        metrics: Dict chứa ít nhất accuracy, precision, recall, f1_score, confusion_matrix
        require_standard: Nếu True thì thiếu key chuẩn sẽ raise ValueError

    Returns:
        Dict sẵn sàng gửi POST /metrics/

    Raises:
        ValueError: Thiếu key chuẩn hoặc giá trị ngoài [0, 1]
    """
    payload = {
        "model_id": model_id,
        "dataset_split_id": dataset_split_id,
        "experiment_run_id": experiment_run_id,
    }

    for key in STANDARD_METRIC_KEYS:
        val = metrics.get(key)
        if key == "confusion_matrix":
            payload[key] = _ensure_confusion_matrix_list(val)
            if require_standard and payload[key] is None:
                raise ValueError(
                    "Chuẩn metrics bắt buộc có 'confusion_matrix' (2D list int hoặc numpy array)"
                )
            continue
        if val is None:
            if require_standard:
                raise ValueError(f"Chuẩn metrics bắt buộc có '{key}'")
            continue
        v = float(val)
        if not (0 <= v <= 1):
            raise ValueError(f"Metric '{key}' phải trong [0, 1], nhận được {v}")
        payload[key] = v

    for key in OPTIONAL_METRIC_KEYS:
        val = metrics.get(key)
        if val is None:
            continue
        if key in ("roc_curve", "pr_curve"):
            payload[key] = val  # dict/list format
            continue
        v = float(val)
        if 0 <= v <= 1:
            payload[key] = v

    return payload


def metrics_from_sklearn(
    y_true: Union[np.ndarray, List[int]],
    y_pred: Union[np.ndarray, List[int]],
    model_id: int,
    dataset_split_id: int,
    experiment_run_id: int,
    *,
    average: str = "weighted",
    labels: Optional[List[int]] = None,
) -> Dict[str, Any]:
    """
    Tính metrics từ sklearn (accuracy, precision, recall, f1, confusion_matrix)
    và trả về payload chuẩn để POST /api/v1/metrics.

    Args:
        y_true: Nhãn thật
        y_pred: Nhãn dự đoán
        model_id: ID model
        dataset_split_id: ID dataset split
        experiment_run_id: ID experiment run
        average: 'weighted', 'macro', 'micro' cho precision/recall/f1
        labels: Danh sách label (optional, dùng cho confusion_matrix)

    Returns:
        Dict dùng cho client.metrics.create(**payload) hoặc POST /metrics/
    """
    try:
        from sklearn.metrics import (
            accuracy_score,
            precision_score,
            recall_score,
            f1_score,
            confusion_matrix,
        )
    except ImportError:
        raise ImportError("Cần cài sklearn: pip install scikit-learn")

    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    accuracy = float(accuracy_score(y_true, y_pred))
    precision = float(precision_score(y_true, y_pred, average=average, zero_division=0))
    recall = float(recall_score(y_true, y_pred, average=average, zero_division=0))
    f1 = float(f1_score(y_true, y_pred, average=average, zero_division=0))
    cm = confusion_matrix(y_true, y_pred, labels=labels)

    metrics = {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1_score": f1,
        "confusion_matrix": cm,
    }
    return build_metrics_payload(
        model_id=model_id,
        dataset_split_id=dataset_split_id,
        experiment_run_id=experiment_run_id,
        metrics=metrics,
        require_standard=True,
    )
