# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, f1_score
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split as sk_train_test_split
from imblearn.over_sampling import SMOTE
from collections import Counter, defaultdict
import os
import re
import gc
import pickle

def load_features_from_npy(train_feat_path, train_lab_path, test_feat_path, test_lab_path, 
                          train_case_ids_path=None, test_case_ids_path=None):
    """
    Loads .npy files produced by the feature extraction script.
    """
    print("="*60)
    print("📂 LOADING DATA")
    print("="*60)
    
    print(f"\nTraining features: {train_feat_path}")
    X_train = np.load(train_feat_path)
    print(f"Training labels: {train_lab_path}")
    y_train = np.load(train_lab_path)
    
    print(f"\nTest features: {test_feat_path}")
    X_test = np.load(test_feat_path)
    print(f"Test labels: {test_lab_path}")
    y_test = np.load(test_lab_path)
    
    print(f"\n✅ Data loaded:")
    print(f"  Training set: {X_train.shape}")
    print(f"  Test set: {X_test.shape}")
    
    # Check label distribution
    print(f"\n📊 Training set label distribution:")
    unique, counts = np.unique(y_train, return_counts=True)
    for u, c in zip(unique, counts):
        print(f"  Class {u}: {c} samples ({c/len(y_train)*100:.1f}%)")
    
    print(f"\n📊 Test set label distribution:")
    unique, counts = np.unique(y_test, return_counts=True)
    for u, c in zip(unique, counts):
        print(f"  Class {u}: {c} samples ({c/len(y_test)*100:.1f}%)")
    
    # Feature statistics
    print(f"\n📈 Feature statistics:")
    print(f"  Train - Min: {X_train.min():.4f}, Max: {X_train.max():.4f}, Mean: {X_train.mean():.4f}, Std: {X_train.std():.4f}")
    print(f"  Test  - Min: {X_test.min():.4f}, Max: {X_test.max():.4f}, Mean: {X_test.mean():.4f}, Std: {X_test.std():.4f}")
    
    # Load case IDs (if available)
    train_cases = None
    if train_case_ids_path and os.path.exists(train_case_ids_path):
        print(f"\nLoading train case IDs: {train_case_ids_path}")
        with open(train_case_ids_path, 'rb') as f:
            train_cases = pickle.load(f)
        print(f"✅ {len(train_cases)} train case IDs loaded")
    
    test_cases = None
    if test_case_ids_path and os.path.exists(test_case_ids_path):
        print(f"\nLoading test case IDs: {test_case_ids_path}")
        with open(test_case_ids_path, 'rb') as f:
            test_cases = pickle.load(f)
        print(f"✅ {len(test_cases)} test case IDs loaded")
    
    return X_train, y_train, X_test, y_test, train_cases, test_cases

def load_and_process_features(features_path, image_paths_path, csv_path, is_train=True):
    """
    Loads features and matches them with the CSV.
    is_train: True -> train data, False -> test data.
    """
    print(f"Loading features from {features_path}")
    features = np.load(features_path)
    with open(image_paths_path, 'r') as f:
        image_paths = [line.strip() for line in f.readlines()]
    
    print(f"Loaded {len(features)} embedding vectors with shape {features.shape}")
    
    # Extract TCGA case IDs from image paths
    print("Extracting TCGA case IDs...")
    tcga_cases = []
    for path in image_paths:
        # Extract the TCGA-XX-XXXX-XXX-XX-XXX pattern
        match = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4}-[0-9A-Z]{3}-[0-9A-Z]{2}-[A-Z0-9]{3})', path)
        if match:
            tcga_case = match.group(1)
        else:
            # If the full pattern is not found, try extracting at least the TCGA-XX-XXXX part
            match = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4})', path)
            if match:
                tcga_case = match.group(1)
            else:
                tcga_case = os.path.basename(os.path.dirname(path))  # Get folder name from the file path
        
        tcga_cases.append(tcga_case)
    
    # Load the CSV data
    df = pd.read_csv(csv_path)
    print(f"CSV loaded with {len(df)} rows")
    
    # Filter CSV based on train/test split
    print(f"Filtering CSV for {'train' if is_train else 'test'} data...")
    if is_train:
        # For the train data: files from the dx_tcga_cropped_20x_train folder
        filtered_df = df[df['filename'].str.contains('dx_tcga_cropped_20x_train', na=False)]
    else:
        # For the test data: files from the dx_tcga_cropped_20x_test folder
        filtered_df = df[df['filename'].str.contains('dx_tcga_cropped_20x_test', na=False)]
    
    print(f"Filtered CSV has {len(filtered_df)} rows for {'train' if is_train else 'test'}")
    
    # Create case-to-grade mapping
    print("Creating case-to-grade mapping...")
    case_to_grade = {}
    for idx, row in df.iterrows():
        filename = row['filename']
        grade = row['gleason_grade']
        match = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4})', filename)
        if match:
            case_id = match.group(1)
            case_to_grade[case_id] = grade
    
    print(f"Created mapping for {len(case_to_grade)} unique cases")
    
    # Match embeddings with grades
    print("Matching embeddings with grades...")
    matched_features = []
    matched_labels = []
    matched_cases = []
    
    # Batch processing for memory management
    batch_size = 10000
    num_batches = len(tcga_cases) // batch_size + (1 if len(tcga_cases) % batch_size > 0 else 0)
    
    for batch_idx in range(num_batches):
        start_idx = batch_idx * batch_size
        end_idx = min((batch_idx + 1) * batch_size, len(tcga_cases))
        
        print(f"Processing batch {batch_idx+1}/{num_batches} (samples {start_idx}-{end_idx})...")
        
        batch_features = []
        batch_labels = []
        batch_cases = []
        
        for i in range(start_idx, end_idx):
            case_id = tcga_cases[i]
            
            # Try exact match first
            if case_id in case_to_grade:
                batch_features.append(features[i])
                batch_labels.append(case_to_grade[case_id])
                batch_cases.append(case_id)
            else:
                # Try short version (TCGA-XX-XXXX)
                short_match = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4})', case_id)
                if short_match:
                    short_id = short_match.group(1)
                    if short_id in case_to_grade:
                        batch_features.append(features[i])
                        batch_labels.append(case_to_grade[short_id])
                        batch_cases.append(short_id)
        
        # Add batch to main lists
        matched_features.extend(batch_features)
        matched_labels.extend(batch_labels)
        matched_cases.extend(batch_cases)
        
        # Clean memory
        del batch_features, batch_labels, batch_cases
        gc.collect()
    
    print(f"Total matched samples: {len(matched_features)}")
    print(f"Unique Gleason grades: {np.unique(matched_labels)}")
    
    return np.array(matched_features), np.array(matched_labels), matched_cases

def patient_level_split(X, y, case_ids, test_size=0.15, random_state=42):
    """
    Performs a patient-level train/validation split.
    All samples from the same patient will be assigned to either train or validation.
    This is critical to prevent data leakage.
    """
    if case_ids is None or len(case_ids) == 0:
        print("⚠️ No case IDs found; using stratified random split...")
        return sk_train_test_split(X, y, test_size=test_size, random_state=random_state, stratify=y)
    
    # Group samples for each patient
    patient_to_indices = defaultdict(list)
    for idx, case_id in enumerate(case_ids):
        patient_to_indices[case_id].append(idx)
    
    # List patients
    unique_patients = list(patient_to_indices.keys())
    print(f"\n📋 Total {len(unique_patients)} unique patients found")
    
    # Determine each patient's label (majority vote)
    patient_labels = {}
    for patient_id, indices in patient_to_indices.items():
        patient_labels[patient_id] = Counter(y[indices]).most_common(1)[0][0]
    
    # Patient-based split for stratification
    patient_labels_list = [patient_labels[p] for p in unique_patients]
    train_patients, val_patients = sk_train_test_split(
        unique_patients, 
        test_size=test_size, 
        random_state=random_state,
        stratify=patient_labels_list
    )
    
    # Collect indices
    train_indices = []
    val_indices = []
    
    for patient_id in train_patients:
        train_indices.extend(patient_to_indices[patient_id])
    
    for patient_id in val_patients:
        val_indices.extend(patient_to_indices[patient_id])
    
    train_indices = np.array(train_indices)
    val_indices = np.array(val_indices)
    
    print(f"✅ Patient-level split:")
    print(f"  Training: {len(train_patients)} patients, {len(train_indices)} samples")
    print(f"  Validation: {len(val_patients)} patients, {len(val_indices)} samples")
    
    return X[train_indices], X[val_indices], y[train_indices], y[val_indices]

def main():
    # Create the output directory
    output_dir = os.path.join('evaluation', 'mlp_results')
    os.makedirs(output_dir, exist_ok=True)
    print(f"📁 Results will be saved to: {output_dir}")
    
    # Load files produced by the feature extraction script
    X_train, y_train, X_test, y_test, train_cases, test_cases = load_features_from_npy(
        train_feat_path='features_train_epoch64.npy',
        train_lab_path='labels_train_epoch64.npy',
        test_feat_path='features_test_epoch64.npy',
        test_lab_path='labels_test_epoch64.npy',
        train_case_ids_path='case_ids_train.pkl',
        test_case_ids_path='case_ids_test.pkl'
    )
    
    # Labels may already be numeric; check
    print("\n" + "="*60)
    print("🏷️ LABEL CHECK")
    print("="*60)
    
    # If labels are strings, encode them; otherwise use them as-is
    if y_train.dtype == object or isinstance(y_train[0], str):
        print("Labels are strings; encoding...")
        label_encoder = LabelEncoder()
        y_train_encoded = label_encoder.fit_transform(y_train)
        y_test_encoded = label_encoder.transform(y_test)
        print(f"Unique classes after encoding: {np.unique(y_train_encoded)}")
        print("Label mapping:", dict(zip(label_encoder.classes_, range(len(label_encoder.classes_)))))
    else:
        print("Labels are already in numeric format.")
        # First, copy the labels
        y_train_encoded = y_train.copy()
        y_test_encoded = y_test.copy()
        
        # Use LabelEncoder as well for numeric labels (so it can be pickled)
        unique_labels = np.unique(y_train_encoded)
        label_encoder = LabelEncoder()
        # Store class names as numeric values
        label_encoder.classes_ = unique_labels
        # inverse_transform already works because LabelEncoder supports it
        print(f"Unique classes after encoding: {np.unique(y_train_encoded)}")
        print("Label mapping:", {int(c): int(c) for c in unique_labels})

    # IMPORTANT: Normalize features before applying SMOTE
    print("\n" + "="*60)
    print("📏 NORMALIZING FEATURES (BEFORE SMOTE)")
    print("="*60)
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    print(f"✅ Normalization completed")

    # IMPORTANT: Patient-level train/validation split (before SMOTE)
    print("\n" + "="*60)
    print("👥 PATIENT-LEVEL TRAIN/VALIDATION SPLIT")
    print("="*60)
    print("⚠️ Critical: All patches from the same patient will be in either train or validation!")
    
    if train_cases is not None and len(train_cases) > 0:
        X_train_final, X_val, y_train_final, y_val = patient_level_split(
            X_train_scaled, y_train_encoded, train_cases, test_size=0.15, random_state=42
        )
        print("✅ Patient-level validation split successful")
    else:
        # Otherwise, use stratified random split (not ideal)
        print("⚠️ Train case IDs not found; using stratified random split...")
        X_train_final, X_val, y_train_final, y_val = sk_train_test_split(
            X_train_scaled, y_train_encoded, 
            test_size=0.15, 
            random_state=42, 
            stratify=y_train_encoded
        )
        print(f"✅ Random split: {len(X_train_final)} train, {len(X_val)} validation")

    # Print class distribution before SMOTE
    print("\n" + "="*60)
    print("📊 CLASS DISTRIBUTION (BEFORE SMOTE - TRAIN SET ONLY)")
    print("="*60)
    class_dist_before = Counter(y_train_final)
    for cls, count in sorted(class_dist_before.items()):
        print(f"  Class {cls}: {count} samples ({count/len(y_train_final)*100:.1f}%)")

    # IMPORTANT: Apply SMOTE only to the train set (do not touch the validation set!)
    print("\n" + "="*60)
    print("🔄 APPLYING SMOTE (ONLY TO TRAIN SET - DO NOT APPLY TO VALIDATION)")
    print("="*60)
    print("⚠️ Critical: SMOTE is not applied to the validation set; only the train set!")
    try:
        smote = SMOTE(random_state=42, k_neighbors=min(5, min(class_dist_before.values())-1))
        X_train_resampled, y_train_resampled = smote.fit_resample(X_train_final, y_train_final)
        
        # Print class distribution after SMOTE
        print("\n📊 Class distribution (after SMOTE - Train set):")
        class_dist_after = Counter(y_train_resampled)
        for cls, count in sorted(class_dist_after.items()):
            print(f"  Class {cls}: {count} samples ({count/len(y_train_resampled)*100:.1f}%)")
        print(f"\n✅ SMOTE successful: {len(X_train_final)} -> {len(X_train_resampled)} samples")
    except Exception as e:
        print(f"⚠️ Could not apply SMOTE: {e}")
        print("Continuing without SMOTE...")
        X_train_resampled = X_train_final
        y_train_resampled = y_train_final

    print(f"\n✅ Data preparation completed:")
    print(f"  Train (after SMOTE): {X_train_resampled.shape}")
    print(f"  Validation: {X_val.shape}")
    print(f"  Test: {X_test_scaled.shape}")

    # Create and train MLP model with optimized parameters
    print("\n" + "="*60)
    print("🧠 Training MLP model (Small architecture + strong regularization)")
    print("="*60)
    
    # Select architecture based on feature dimension (smaller!)
    feature_dim = X_train_resampled.shape[1]
    n_classes = len(np.unique(y_train_resampled))
    n_samples = len(X_train_resampled)
    
    print(f"📐 Model parameters:")
    print(f"  Feature boyutu: {feature_dim}")
    print(f"  Number of classes: {n_classes}")
    print(f"  Number of training examples: {n_samples}")
    
    # Optimized architecture for %90+ accuracy
    # Larger but balanced architecture (capacity increase + overfitting control)
    if feature_dim >= 512:
        hidden_layers = (1024, 512, 256)  # Deep network for large features
    elif feature_dim >= 256:
        hidden_layers = (512, 256, 128)   # Medium-sized features
    else:
        hidden_layers = (256, 128, 64)    # Small feature sizes
    
    print(f"  Hidden layers: {hidden_layers} (optimized for %90+ accuracy)")
    
    # Compute class weights (for class imbalance)
    from sklearn.utils.class_weight import compute_sample_weight
    class_weights = compute_sample_weight('balanced', y_train_resampled)
    
    print(f"\n⚖️ Class weights computed (balanced)")
    
    # Try different hyperparameter combinations
    best_mlp = None
    best_val_score = -1
    best_params = None
    
    # Hyperparameter grid
    alpha_values = [0.0001, 0.001, 0.01]  # Regularization
    lr_values = [0.0005, 0.001, 0.002]    # Learning rate
    
    print(f"\n🔍 Hyperparameter tuning starting...")
    print(f"  Alpha values: {alpha_values}")
    print(f"  Learning rate values: {lr_values}")
    
    for alpha in alpha_values:
        for lr in lr_values:
            print(f"\n  Testing: alpha={alpha}, lr={lr}")
            
            mlp_temp = MLPClassifier(
                hidden_layer_sizes=hidden_layers,
                activation='relu',
                solver='adam',
                alpha=alpha,  # Regularization
                batch_size=128,  # Smaller batch size (better gradient)
                learning_rate='adaptive',
                learning_rate_init=lr,
                max_iter=500,  # Max iterations for each configuration
                early_stopping=True,  # Enable early stopping
                validation_fraction=0.1,  # 10% for validation
                n_iter_no_change=20,  # Stop if no improvement for 20 iterations
                tol=1e-5,  # More strict tolerance
                random_state=42,
                verbose=False,  # Silent during tuning
                beta_1=0.9,
                beta_2=0.999,
                epsilon=1e-8
            )
            
            # Train with class weights
            mlp_temp.fit(X_train_resampled, y_train_resampled, sample_weight=class_weights)
            
            # Evaluate on the validation set
            val_score = mlp_temp.score(X_val, y_val)
            print(f"    Validation Score: {val_score:.6f}")
            
            if val_score > best_val_score:
                best_val_score = val_score
                best_mlp = mlp_temp
                best_params = {'alpha': alpha, 'lr': lr}
                print(f"    ✅ New best score!")
    
    print(f"\n✅ Best parameters: {best_params}")
    print(f"✅ Best validation score: {best_val_score:.6f}")
    
    # Retrain the best model on the full train set (more iterations)
    print(f"\n🎯 Training final model (with best parameters)...")
    mlp = MLPClassifier(
        hidden_layer_sizes=hidden_layers,
        activation='relu',
        solver='adam',
        alpha=best_params['alpha'],
        batch_size=128,
        learning_rate='adaptive',
        learning_rate_init=best_params['lr'],
        max_iter=2000,  # More iterations for final training
        early_stopping=True,
        validation_fraction=0.1,
        n_iter_no_change=30,
        tol=1e-5,
        random_state=42,
        verbose=True,
        beta_1=0.9,
        beta_2=0.999,
        epsilon=1e-8
    )
    
    # Final training with class weights
    mlp.fit(X_train_resampled, y_train_resampled, sample_weight=class_weights)
    
    # Evaluate on the validation set
    val_score = mlp.score(X_val, y_val)
    print(f"\n📊 Final Validation Score: {val_score:.6f}")

    # Evaluate on test set
    print("\n" + "="*60)
    print("📊 Evaluation on test set")
    print("="*60)
    y_pred = mlp.predict(X_test_scaled)
    y_pred_proba = mlp.predict_proba(X_test_scaled)
    
    # Accuracy ve F1 hesapla
    acc = accuracy_score(y_test_encoded, y_pred)
    f1 = f1_score(y_test_encoded, y_pred, average='weighted')
    f1_macro = f1_score(y_test_encoded, y_pred, average='macro')
    
    print(f"\n🎯 Genel Metrikler:")
    print(f"  Accuracy: {acc:.4f} ({acc*100:.2f}%)")
    print(f"  F1-Score (weighted): {f1:.4f}")
    print(f"  F1-Score (macro): {f1_macro:.4f}")
    
    # Convert numeric predictions back to original labels for the report
    if hasattr(label_encoder, 'inverse_transform') and callable(label_encoder.inverse_transform):
        try:
            y_test_original = label_encoder.inverse_transform(y_test_encoded)
            y_pred_original = label_encoder.inverse_transform(y_pred)
        except:
            y_test_original = y_test_encoded
            y_pred_original = y_pred
    else:
        y_test_original = y_test_encoded
        y_pred_original = y_pred
    
    print("\n📋 Detailed Classification Report:")
    print(classification_report(y_test_original, y_pred_original, digits=4))

    # Create confusion matrix
    cm = confusion_matrix(y_test_original, y_pred_original)
    plt.figure(figsize=(12, 10))
    
    # Prepare class names
    if hasattr(label_encoder, 'classes_'):
        class_names = [str(c) for c in label_encoder.classes_]
    else:
        unique_classes = sorted(np.unique(y_test_original))
        class_names = [f'Class_{c}' for c in unique_classes]
    
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=class_names,
                yticklabels=class_names,
                cbar_kws={'label': 'Number of Samples'})
    plt.xlabel('Predicted Label', fontsize=12)
    plt.ylabel('True Label', fontsize=12)
    plt.title(f'MLP Confusion Matrix (Accuracy: {acc:.4f})', fontsize=14)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    plt.tight_layout()
    confusion_matrix_path = os.path.join(output_dir, 'mlp_confusion_matrix_fused.png')
    plt.savefig(confusion_matrix_path, dpi=300, bbox_inches='tight')
    print(f"✅ Confusion matrix saved: {confusion_matrix_path}")

    # Save the trained model, scaler and label encoder
    print("\n" + "="*60)
    print("💾 SAVING MODEL")
    print("="*60)
    
    # Print training information
    print(f"\n📊 Training Summary:")
    print(f"  Toplam iterasyon: {mlp.n_iter_}")
    print(f"  Final loss: {mlp.loss_curve_[-1]:.6f}" if hasattr(mlp, 'loss_curve_') and mlp.loss_curve_ else "  Final loss: N/A")
    if hasattr(mlp, 'validation_scores_') and mlp.validation_scores_:
        print(f"  Final validation score: {mlp.validation_scores_[-1]:.6f}")
    
    model_path = os.path.join(output_dir, 'mlp_model_fused.joblib')
    scaler_path = os.path.join(output_dir, 'mlp_scaler_fused.joblib')
    
    joblib.dump(mlp, model_path)
    joblib.dump(scaler, scaler_path)
    
    # Save the label encoder (it is always a LabelEncoder instance and is pickle-able)
    encoder_path = os.path.join(output_dir, 'mlp_label_encoder_fused.joblib')
    if label_encoder is not None:
        try:
            joblib.dump(label_encoder, encoder_path)
            print(f"✅ Label encoder saved: {encoder_path}")
        except Exception as e:
            print(f"⚠️ Could not save the label encoder: {e}")
            print("   Saving class mapping manually...")
            # Alternative: save class mapping as a dict
            class_mapping = {
                'classes_': label_encoder.classes_.tolist() if hasattr(label_encoder, 'classes_') else None,
                'type': 'numeric' if y_train.dtype != object else 'string'
            }
            import json
            mapping_path = os.path.join(output_dir, 'mlp_label_encoder_mapping.json')
            with open(mapping_path, 'w') as f:
                json.dump(class_mapping, f)
            print(f"✅ Class mapping saved: {mapping_path}")
    
    print(f"✅ Model saved: {model_path}")
    print(f"✅ Scaler saved: {scaler_path}")
    
    # Save patient-level results (if case IDs are available)
    if test_cases is not None:
        results_df = pd.DataFrame({
            'case_id': test_cases,
            'true_label': y_test_original,
            'pred_label': y_pred_original,
            'correct': (y_test_encoded == y_pred).astype(int),
            'confidence': np.max(y_pred_proba, axis=1)
        })
        results_path = os.path.join(output_dir, 'mlp_patient_level_results.csv')
        results_df.to_csv(results_path, index=False)
        print(f"✅ Patient-level results saved: {results_path}")

    # Plot learning curve
    if hasattr(mlp, 'loss_curve_') and mlp.loss_curve_ is not None:
        plt.figure(figsize=(12, 8))
        plt.plot(mlp.loss_curve_)
        plt.title('MLP Learning Curve - Fused Features', fontsize=14)
        plt.xlabel('Iterations', fontsize=12)
        plt.ylabel('Loss', fontsize=12)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        learning_curve_path = os.path.join(output_dir, 'mlp_learning_curve_fused.png')
        plt.savefig(learning_curve_path, dpi=300, bbox_inches='tight')
        print(f"✅ Learning curve saved: {learning_curve_path}")

    # Validation score curve
    if hasattr(mlp, 'validation_scores_') and mlp.validation_scores_ is not None:
        plt.figure(figsize=(12, 8))
        plt.plot(mlp.validation_scores_)
        plt.title('MLP Validation Score Curve - Fused Features', fontsize=14)
        plt.xlabel('Iterations', fontsize=12)
        plt.ylabel('Validation Score', fontsize=12)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        validation_curve_path = os.path.join(output_dir, 'mlp_validation_curve_fused.png')
        plt.savefig(validation_curve_path, dpi=300, bbox_inches='tight')
        print(f"✅ Validation curve saved: {validation_curve_path}")

    print("\n" + "="*60)
    print("🎉 PROCESS COMPLETED!")
    print("="*60)
    print("\nThe model is now ready to make predictions.")

# Function to predict on new samples
def predict_gleason_grade(embedding_vector, 
                         model_path=os.path.join('evaluation', 'mlp_results', 'mlp_model_fused.joblib'),
                         scaler_path=os.path.join('evaluation', 'mlp_results', 'mlp_scaler_fused.joblib'),
                         encoder_path=os.path.join('evaluation', 'mlp_results', 'mlp_label_encoder_fused.joblib')):
    """Predict Gleason grade for a new DINO embedding vector"""
    model = joblib.load(model_path)
    scaler = joblib.load(scaler_path)
    label_encoder = joblib.load(encoder_path)
    
    # Reshape and scale the input
    embedding_vector = np.array(embedding_vector).reshape(1, -1)
    embedding_vector_scaled = scaler.transform(embedding_vector)
    
    # Get prediction and probabilities
    prediction = model.predict(embedding_vector_scaled)
    probabilities = model.predict_proba(embedding_vector_scaled)
    
    # Convert numeric prediction back to original label
    prediction_original = label_encoder.inverse_transform(prediction)
    
    return {
        'predicted_grade': prediction_original[0],
        'probabilities': dict(zip(label_encoder.classes_, probabilities[0]))
    }

if __name__ == "__main__":
    main()