"""
Model evaluation utilities
"""

import numpy as np
from sklearn.metrics import (
    accuracy_score, f1_score, recall_score, precision_score,
    roc_curve, auc, average_precision_score, confusion_matrix,
    roc_auc_score, precision_recall_curve
)
from typing import Dict, Any, Optional, Tuple
from tensorflow import keras
from kwslib.utils.metrics_template import MetricsTemplate


class ModelEvaluator:
    """Helper class for evaluating models"""
    
    def __init__(self, api_client=None, use_mlflow: bool = False, mlflow_run_id: Optional[str] = None):
        """
        Initialize model evaluator.
        
        Args:
            api_client: Optional API client for MLFlow logging
            use_mlflow: Whether to use MLFlow logging
            mlflow_run_id: MLFlow run ID if using MLFlow
        """
        self.api = api_client
        self.use_mlflow = use_mlflow
        self.mlflow_run_id = mlflow_run_id
    
    def evaluate(
        self,
        model: keras.Model,
        X_test: np.ndarray,
        y_test: np.ndarray,
        label_to_id: Optional[Dict[str, int]] = None
    ) -> Tuple[Dict[str, Any], np.ndarray, np.ndarray]:
        """
        Evaluate model on test set with full metrics including ROC, AUC, Confusion Matrix.
        
        Args:
            model: Trained Keras model
            X_test: Test features
            y_test: Test labels (integer encoded)
            label_to_id: Optional mapping from label name to integer ID
        
        Returns:
            Tuple of (metrics_dict, y_pred, y_pred_proba)
            - metrics_dict: Full metrics dictionary following MetricsTemplate
            - y_pred: Predicted labels
            - y_pred_proba: Prediction probabilities
        """
        print("\n=== Evaluating on test set ===")
        
        # Evaluate using Keras
        test_loss, test_accuracy = model.evaluate(X_test, y_test, verbose=0)
        print(f"Test Loss: {test_loss:.4f}")
        print(f"Test Accuracy: {test_accuracy:.4f}")
        
        # Get predictions
        y_pred_proba = model.predict(X_test, verbose=0)
        y_pred = np.argmax(y_pred_proba, axis=1)
        
        # Calculate basic metrics
        accuracy = accuracy_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred, average='weighted')
        recall = recall_score(y_test, y_pred, average='weighted')
        precision = precision_score(y_test, y_pred, average='weighted')
        
        # Initialize metrics template
        metrics = MetricsTemplate.create_empty()
        metrics['accuracy'] = float(accuracy)
        metrics['f1_score'] = float(f1)
        metrics['recall'] = float(recall)
        metrics['precision'] = float(precision)
        
        # Calculate ROC and AUC (for binary and multiclass)
        num_classes = len(np.unique(y_test))
        
        if num_classes == 2:
            # Binary classification
            y_test_binary = (y_test == 1).astype(int)
            y_pred_proba_binary = y_pred_proba[:, 1]
            
            # ROC curve
            fpr, tpr, _ = roc_curve(y_test_binary, y_pred_proba_binary)
            metrics['roc_curve'] = MetricsTemplate.format_roc_curve(fpr, tpr)
            
            # AUC
            auc_macro = roc_auc_score(y_test_binary, y_pred_proba_binary)
            metrics['auc_macro'] = float(auc_macro)
            metrics['auc_micro'] = float(auc_macro)  # Same for binary
            
            # PR curve
            precision_curve, recall_curve, _ = precision_recall_curve(y_test_binary, y_pred_proba_binary)
            metrics['pr_curve'] = MetricsTemplate.format_pr_curve(recall_curve, precision_curve)
            
            # Average Precision
            ap_macro = average_precision_score(y_test_binary, y_pred_proba_binary)
            metrics['ap_macro'] = float(ap_macro)
            metrics['ap_micro'] = float(ap_macro)
            
        else:
            # Multiclass classification
            # One-vs-rest ROC for each class
            fpr_dict = {}
            tpr_dict = {}
            auc_scores = []
            
            for i in range(num_classes):
                y_test_binary = (y_test == i).astype(int)
                y_pred_proba_binary = y_pred_proba[:, i]
                
                if len(np.unique(y_test_binary)) > 1:  # Check if class exists in test set
                    fpr_i, tpr_i, _ = roc_curve(y_test_binary, y_pred_proba_binary)
                    fpr_dict[i] = fpr_i
                    tpr_dict[i] = tpr_i
                    auc_i = roc_auc_score(y_test_binary, y_pred_proba_binary)
                    auc_scores.append(auc_i)
            
            # Use macro-averaged ROC (average across classes)
            if fpr_dict:
                # Interpolate to common FPR values
                all_fpr = np.unique(np.concatenate([fpr_dict[i] for i in fpr_dict.keys()]))
                mean_tpr = np.zeros_like(all_fpr)
                
                for i in fpr_dict.keys():
                    mean_tpr += np.interp(all_fpr, fpr_dict[i], tpr_dict[i])
                
                mean_tpr /= len(fpr_dict)
                metrics['roc_curve'] = MetricsTemplate.format_roc_curve(all_fpr, mean_tpr)
            
            # AUC macro and micro
            if auc_scores:
                metrics['auc_macro'] = float(np.mean(auc_scores))
                metrics['auc_micro'] = float(np.mean(auc_scores))  # Simplified for multiclass
            
            # PR curve (macro-averaged)
            precision_curves = []
            recall_curves = []
            ap_scores = []
            
            for i in range(num_classes):
                y_test_binary = (y_test == i).astype(int)
                y_pred_proba_binary = y_pred_proba[:, i]
                
                if len(np.unique(y_test_binary)) > 1:
                    precision_i, recall_i, _ = precision_recall_curve(y_test_binary, y_pred_proba_binary)
                    precision_curves.append(precision_i)
                    recall_curves.append(recall_i)
                    ap_i = average_precision_score(y_test_binary, y_pred_proba_binary)
                    ap_scores.append(ap_i)
            
            if precision_curves:
                # Average PR curves: lưới recall cố định, interpolate từng curve rồi trung bình
                try:
                    recall_grid = np.linspace(0, 1, 101)
                    mean_precision = np.zeros_like(recall_grid, dtype=float)
                    n_curves = 0
                    for i in range(len(precision_curves)):
                        rec_i = np.asarray(recall_curves[i], dtype=float)
                        prec_i = np.asarray(precision_curves[i], dtype=float)
                        if len(rec_i) != len(prec_i) or len(rec_i) < 2:
                            continue
                        sort_idx = np.argsort(rec_i)
                        rec_i = rec_i[sort_idx]
                        prec_i = prec_i[sort_idx]
                        interp_prec = np.interp(recall_grid, rec_i, prec_i)
                        mean_precision += interp_prec
                        n_curves += 1
                    if n_curves > 0:
                        mean_precision /= n_curves
                        metrics['pr_curve'] = MetricsTemplate.format_pr_curve(recall_grid, mean_precision)
                except Exception:
                    pass  # bỏ qua pr_curve nếu lỗi, vẫn giữ accuracy/precision/recall/f1/cm
                if ap_scores:
                    metrics['ap_macro'] = float(np.mean(ap_scores))
                    metrics['ap_micro'] = float(np.mean(ap_scores))
        
        # Confusion Matrix
        cm = confusion_matrix(y_test, y_pred)
        metrics['confusion_matrix'] = MetricsTemplate.format_confusion_matrix(cm)
        
        print(f"\n=== Evaluation Metrics ===")
        print(f"Accuracy:  {metrics['accuracy']:.4f}")
        print(f"F1 Score:  {metrics['f1_score']:.4f}")
        print(f"Recall:    {metrics['recall']:.4f}")
        print(f"Precision: {metrics['precision']:.4f}")
        if metrics['auc_macro'] is not None:
            print(f"AUC Macro: {metrics['auc_macro']:.4f}")
        if metrics['auc_micro'] is not None:
            print(f"AUC Micro: {metrics['auc_micro']:.4f}")
        
        # Log to MLFlow if enabled
        if self.use_mlflow and self.mlflow_run_id and self.api:
            try:
                mlflow_metrics = {
                    'test_loss': float(test_loss),
                    'test_accuracy': float(accuracy),
                    'test_f1_score': float(f1),
                    'test_recall': float(recall),
                    'test_precision': float(precision),
                }
                if metrics['auc_macro'] is not None:
                    mlflow_metrics['test_auc_macro'] = metrics['auc_macro']
                if metrics['auc_micro'] is not None:
                    mlflow_metrics['test_auc_micro'] = metrics['auc_micro']
                
                self.api.mlflow.log_metrics(
                    run_id=self.mlflow_run_id,
                    metrics=mlflow_metrics
                )
                print("Logged final metrics to MLFlow")
                
                # Update run status to FINISHED
                self.api.mlflow.update_run_status(
                    run_id=self.mlflow_run_id,
                    status="FINISHED"
                )
                print("Updated MLFlow run status to FINISHED")
            except Exception as e:
                print(f"Warning: Failed to log final metrics to MLFlow: {e}")
        
        return metrics, y_pred, y_pred_proba
