# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/01_report.ipynb.

# %% auto 0
__all__ = ['run_gnb_optimization_and_evaluation']

# %% ../notebooks/01_report.ipynb 88
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import precision_recall_curve

def run_gnb_optimization_and_evaluation(X_train, y_train, X_test, y_test, target_recall=0.91):
    """
    Performs grid search on smoothing and class weights, 
    visualizes the heatmap, and evaluates the final surgical model.
    """
    
    # 1. Identify 'hard weed' indices (Weeds often misclassified as crops)
    # This ensures the function is self-contained
    baseline_gnb = GaussianNB().fit(X_train, y_train)
    baseline_preds = baseline_gnb.predict(X_train)
    # Hard weeds are ground-truth 'weed' but predicted as 'crop'
    hard_weed_indices = np.where((y_train == "weed") & (baseline_preds == "crop"))[0]

    # 2. Define the grid
    smoothing_grid = np.logspace(-9, 1, 10)
    weight_grid = np.linspace(1.0, 5.0, 10)
    auc_results = np.zeros((len(smoothing_grid), len(weight_grid)))

    best_auc = 0
    best_params = {"smooth": 0, "weight": 0}

    print("Running Grid Search...")
    for i, s in enumerate(smoothing_grid):
        for j, w in enumerate(weight_grid):
            iter_weights = np.ones(len(y_train))
            iter_weights[hard_weed_indices] = w
            
            # Using your existing experiment function
            m, _, _ = experiment_gaussian_nb(X_train, y_train, X_test, y_test, 
                                           var_smoothing=s, 
                                           sample_weight=iter_weights)
            
            auc_results[i, j] = m["AUC score"]
            if m["AUC score"] > best_auc:
                best_auc = m["AUC score"]
                best_params = {"smooth": s, "weight": w}

    # 3. Plot Heatmap
    plt.figure(figsize=(10, 8))
    sns.heatmap(auc_results, annot=True, fmt=".3f", 
                xticklabels=[f"{w:.1f}" for w in weight_grid], 
                yticklabels=[f"{s:.1e}" for s in smoothing_grid],
                cmap="YlGnBu")
    plt.title("Grid Search: AUC Score (Smoothing vs. Weed Weight)")
    plt.xlabel("Hard Weed Weight")
    plt.ylabel("Var Smoothing")
    plt.show()

    print(f"Optimal found: Smoothing={best_params['smooth']:.2e}, Weight={best_params['weight']:.2f}")

    # 4. Final "Surgical" Model Training
    final_weights = np.ones(len(y_train))
    final_weights[hard_weed_indices] = best_params['weight']
    
    final_nb_clf = GaussianNB(var_smoothing=best_params['smooth'])
    final_nb_clf.fit(X_train, y_train, sample_weight=final_weights)
    
    # Probabilities for 'crop' (Column 0)
    final_probs = final_nb_clf.predict_proba(X_test)[:, 0]

    # 5. Threshold Pinning for High Recall
    precisions, recalls, thresholds = precision_recall_curve(y_test, final_probs, pos_label="crop")
    
    # Find the threshold for target_recall
    idx = np.where(recalls >= target_recall)[0][-1]
    pinned_t = thresholds[idx]

    # 6. Final Evaluation
    final_preds = ["crop" if p >= pinned_t else "weed" for p in final_probs]
    print(f"\n--- FINAL MODEL PERFORMANCE (Pinned Threshold: {pinned_t:.4f}) ---")
    results = evaluate(y_test, final_preds, y_probs=final_probs, print_cnfm=True)
    
    return final_nb_clf, pinned_t