# -*- coding: utf-8 -*-
"""
Enhanced Logistic Regression Classifier for Gleason Grade Classification
- Patient-level GroupKFold cross-validation (grouped by `case_id`)
- Advanced preprocessing: PCA, feature selection, different normalizations
- Extended hyperparameter tuning (C, solver, penalty)
- Ensemble approach: combining multiple models
- SMOTE for class imbalance handling
- Results are saved to `evaluation/logistic_regression_results`
"""
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GroupKFold
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, f1_score
from sklearn.decomposition import PCA
from sklearn.feature_selection import SelectKBest, f_classif, mutual_info_classif
from sklearn.preprocessing import StandardScaler, LabelEncoder, MinMaxScaler, RobustScaler
from sklearn.preprocessing import normalize
from sklearn.ensemble import VotingClassifier
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
from collections import Counter, defaultdict
from imblearn.over_sampling import SMOTE
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 main():
    # Create output directory
    output_dir = os.path.join('evaluation', 'logistic_regression_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 for numeric labels too (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
        print(f"Unique classes after encoding: {np.unique(y_train_encoded)}")
        print("Label mapping:", {int(c): int(c) for c in unique_labels})

    # ADVANCED PREPROCESSING
    print("\n" + "="*60)
    print("📏 ADVANCED FEATURE PROCESSING")
    print("="*60)
    
    # First, apply basic normalization
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    print(f"✅ StandardScaler normalization completed")
    
    # Feature selection: keep the most important features
    # NOTE: To avoid overly aggressive feature selection, keep more features
    print("\n🔍 Applying Feature Selection...")
    original_n_features = X_train_scaled.shape[1]
    
    # Select based on the number of features
    if original_n_features > 1000:
        n_features_to_select = min(500, int(original_n_features * 0.7))  # %70'ini tut
    elif original_n_features > 500:
        n_features_to_select = min(400, int(original_n_features * 0.8))  # %80'ini tut
    else:
        n_features_to_select = int(original_n_features * 0.9)  # keep 90%
    
    # Minimum number of features is guaranteed
    n_features_to_select = max(n_features_to_select, 50)
    
    selector = SelectKBest(score_func=f_classif, k=n_features_to_select)
    X_train_selected = selector.fit_transform(X_train_scaled, y_train_encoded)
    X_test_selected = selector.transform(X_test_scaled)
    print(f"✅ Feature selection: {X_train_scaled.shape[1]} -> {X_train_selected.shape[1]} feature")
    
    # PCA dimensionality reduction (optional - if the number of features is still large)
    if X_train_selected.shape[1] > 500:
        print("\n📊 Applying PCA...")
        pca = PCA(n_components=0.95)  # Keep 95% of the variance
        X_train_pca = pca.fit_transform(X_train_selected)
        X_test_pca = pca.transform(X_test_selected)
        print(f"✅ PCA: {X_train_selected.shape[1]} -> {X_train_pca.shape[1]} dimensions")
        X_train_final = X_train_pca
        X_test_final = X_test_pca
    else:
        X_train_final = X_train_selected
        X_test_final = X_test_selected
        pca = None
    
    print(f"✅ Final feature boyutu: {X_train_final.shape[1]}")

    # Print train set distribution
    print("\n" + "="*60)
    print("📊 TRAIN SET CLASS DISTRIBUTION (BEFORE SMOTE)")
    print("="*60)
    class_dist_before = Counter(y_train_encoded)
    for cls, count in sorted(class_dist_before.items()):
        print(f"  Class {cls}: {count} samples ({count/len(y_train_encoded)*100:.1f}%)")
    
    # Handle class imbalance with SMOTE
    print("\n" + "="*60)
    print("🔄 APPLYING SMOTE")
    print("="*60)
    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_encoded)
        
        print("\n📊 Class distribution (after SMOTE):")
        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_encoded

    print(f"\n✅ Data preparation completed:")
    print(f"  Train (after SMOTE): {X_train_resampled.shape}")
    print(f"  Test : {X_test_final.shape}")
    
    # Create and train Logistic Regression model
    print("\n" + "="*60)
    print("🧠 TRAINING LOGISTIC REGRESSION MODEL")
    print("="*60)
    
    feature_dim = X_train_final.shape[1]
    n_classes = len(np.unique(y_train_resampled))
    n_samples = len(X_train_resampled)
    
    print(f"📐 Model parameters:")
    print(f"  Original feature dimension: {X_train.shape[1]}")
    print(f"  Final feature dimension (selection+PCA): {feature_dim}")
    print(f"  Number of classes: {n_classes}")
    print(f"  Number of training examples (after SMOTE): {n_samples}")
    
    # Logistic Regression parametreleri
    # C: Regularization strength (smaller C = stronger regularization)
    # class_weight: to handle class imbalance
    # max_iter: maximum iteration count
    # solver: 'lbfgs' is often good for small-to-medium datasets
    #         'saga' is better for larger datasets
    
    # EXTENDED HYPERPARAMETER TUNING
    # NOTE: case IDs cannot be preserved after SMOTE (synthetic examples are created)
    # Therefore we use regular KFold (instead of GroupKFold)
    from sklearn.model_selection import KFold
    
    # Assign a unique ID to each example after SMOTE (for GroupKFold-like grouping)
    # This way each example belongs to its own group (prevents data leakage)
    groups = np.arange(len(y_train_resampled))
    print("\n👥 After SMOTE, each example belongs to its own group (prevents data leakage)")
    print("   Using standard KFold (instead of GroupKFold)")
    
    kf = KFold(n_splits=5, shuffle=True, random_state=42)
    
    # Extended hyperparameter grid
    C_values = [0.01, 0.1, 0.5, 1.0, 5.0, 10.0, 50.0, 100.0, 500.0]
    solvers = ['lbfgs', 'saga']  # saga is good for larger datasets
    penalties = ['l2', 'elasticnet']  # elasticnet is more flexible
    
    best_C = 1.0
    best_solver = 'lbfgs'
    best_penalty = 'l2'
    best_cv_score = -1.0
    best_model = None
    best_params = {}

    print(f"\n🔍 Extended Hyperparameter Tuning starting...")
    print(f"  C values: {C_values}")
    print(f"  Solvers: {solvers}")
    print(f"  Penalties: {penalties}")
    
    total_combinations = len(C_values) * len(solvers) * len(penalties)
    current_combination = 0
    
    for C in C_values:
        for solver in solvers:
            for penalty in penalties:
                # Check solver-penalty compatibility
                # lbfgs supports only l2, saga supports elasticnet
                if solver == 'lbfgs' and penalty == 'elasticnet':
                    print(f"\n  ⏭️  Skipping: {solver} does not support {penalty} penalty")
                    continue
                
                current_combination += 1
                print(f"\n  [{current_combination}/{total_combinations}] C={C}, solver={solver}, penalty={penalty} testing...")
                
                # l1_ratio is required for elasticnet
                if penalty == 'elasticnet':
                    l1_ratios = [0.1, 0.5, 0.9]
                else:
                    l1_ratios = [None]
                
                for l1_ratio in l1_ratios:
                    fold_scores = []
                    
                    try:
                        fold_scores = []
                        for fold_idx, (train_idx, val_idx) in enumerate(kf.split(X_train_resampled, y_train_resampled), start=1):
                            X_tr, X_val = X_train_resampled[train_idx], X_train_resampled[val_idx]
                            y_tr, y_val = y_train_resampled[train_idx], y_train_resampled[val_idx]

                            # Removed multi_class parameter (deprecated; default multinomial)
                            lr_temp = LogisticRegression(
                                C=C,
                                class_weight='balanced',
                                max_iter=2000,  # More iterations
                                random_state=42,
                                solver=solver,
                                penalty=penalty,
                                l1_ratio=l1_ratio if penalty == 'elasticnet' else None,
                                n_jobs=-1,
                                verbose=0
                            )

                            lr_temp.fit(X_tr, y_tr)
                            fold_score = lr_temp.score(X_val, y_val)
                            fold_scores.append(fold_score)
                        
                        if len(fold_scores) > 0:
                            mean_score = float(np.mean(fold_scores))
                            std_score = float(np.std(fold_scores))
                            param_key = f"C={C}, solver={solver}, penalty={penalty}"
                            if l1_ratio is not None:
                                param_key += f", l1_ratio={l1_ratio}"
                            best_params[param_key] = mean_score
                            print(f"    -> Average CV score: {mean_score:.6f} (+/- {std_score:.6f})")

                            if mean_score > best_cv_score:
                                best_cv_score = mean_score
                                best_C = C
                                best_solver = solver
                                best_penalty = penalty
                                # Save the model from the last fold
                                best_model = lr_temp
                                print(f"    ✅ New best score!")
                    except Exception as e:
                        print(f"    ⚠️ Error: {str(e)[:100]}")  # Shortened error message
                        continue

    print(f"\n✅ Best parameters:")
    print(f"  C: {best_C}")
    print(f"  Solver: {best_solver}")
    print(f"  Penalty: {best_penalty}")
    print(f"  KFold CV Score: {best_cv_score:.6f}")
    
    # ENSEMBLE APPROACH: combine multiple strong models
    print(f"\n" + "="*60)
    print("🎯 BUILDING ENSEMBLE MODEL")
    print("="*60)
    
    # Select the best 3-5 models
    sorted_params = sorted(best_params.items(), key=lambda x: x[1], reverse=True)
    top_models = sorted_params[:min(5, len(sorted_params))]
    
    # If no model succeeded, use the default model
    if len(top_models) == 0 or best_cv_score < 0:
        print("⚠️ No model succeeded; creating the model with default parameters...")
        lr = LogisticRegression(
            C=1.0,
            class_weight='balanced',
            max_iter=2000,
            random_state=42,
            solver='lbfgs',
            penalty='l2',
            n_jobs=-1,
            verbose=1
        )
        lr.fit(X_train_resampled, y_train_resampled)
        print(f"✅ Default model trained")
    elif len(top_models) == 1:
        # If only one model is available, no ensemble is needed
        print(f"✅ Using a single model (no ensemble needed)")
        lr = LogisticRegression(
            C=best_C,
            class_weight='balanced',
            max_iter=2000,
            random_state=42,
            solver=best_solver,
            penalty=best_penalty,
            n_jobs=-1,
            verbose=1
        )
        lr.fit(X_train_resampled, y_train_resampled)
    else:
        # If multiple models are available, build the ensemble
        print(f"\nBest {len(top_models)} models selected:")
        ensemble_models = []
        
        for param_str, score in top_models:
            print(f"  {param_str}: {score:.6f}")
            # Parametreleri parse et
            params = {}
            for part in param_str.split(', '):
                if '=' in part:
                    key, value = part.split('=')
                    if key == 'C':
                        params['C'] = float(value)
                    elif key == 'solver':
                        params['solver'] = value
                    elif key == 'penalty':
                        params['penalty'] = value
                    elif key == 'l1_ratio':
                        params['l1_ratio'] = float(value)
            
            lr_ensemble = LogisticRegression(
                C=params.get('C', best_C),
                class_weight='balanced',
                max_iter=2000,
                random_state=42,
                solver=params.get('solver', best_solver),
                penalty=params.get('penalty', best_penalty),
                l1_ratio=params.get('l1_ratio', None),
                n_jobs=-1,
                verbose=0
            )
            lr_ensemble.fit(X_train_resampled, y_train_resampled)
            ensemble_models.append(('lr_' + str(len(ensemble_models)), lr_ensemble))
        
        # Create a VotingClassifier
        voting_clf = VotingClassifier(
            estimators=ensemble_models,
            voting='soft',  # Probability-based voting
            n_jobs=-1
        )
        voting_clf.fit(X_train_resampled, y_train_resampled)
        
        # Final model olarak ensemble kullan
        lr = voting_clf
        print(f"✅ Ensemble model built ({len(ensemble_models)} models combined)")

    # Evaluate on test set
    print("\n" + "="*60)
    print("📊 EVALUATION ON TEST SET")
    print("="*60)
    y_pred = lr.predict(X_test_final)
    y_pred_proba = lr.predict_proba(X_test_final)
    
    # 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'Logistic Regression 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, 'logistic_regression_confusion_matrix.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)
    
    model_path = os.path.join(output_dir, 'logistic_regression_model.joblib')
    scaler_path = os.path.join(output_dir, 'logistic_regression_scaler.joblib')
    encoder_path = os.path.join(output_dir, 'logistic_regression_label_encoder.joblib')
    selector_path = os.path.join(output_dir, 'logistic_regression_selector.joblib')
    pca_path = os.path.join(output_dir, 'logistic_regression_pca.joblib')
    
    joblib.dump(lr, model_path)
    joblib.dump(scaler, scaler_path)
    joblib.dump(selector, selector_path)
    if pca is not None:
        joblib.dump(pca, pca_path)
    
    # Save label encoder
    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}")
            import json
            class_mapping = {
                'classes_': label_encoder.classes_.tolist() if hasattr(label_encoder, 'classes_') else None,
                'type': 'numeric' if y_train.dtype != object else 'string'
            }
            mapping_path = os.path.join(output_dir, 'logistic_regression_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}")
    print(f"✅ Feature selector saved: {selector_path}")
    if pca is not None:
        print(f"✅ PCA saved: {pca_path}")
    
    # Save model parameters
    model_info = {
        'best_C': float(best_C),
        'best_solver': best_solver,
        'best_penalty': best_penalty,
        'class_weight': 'balanced',
        'cv_score_kfold': float(best_cv_score),
        'test_accuracy': float(acc),
        'test_f1_weighted': float(f1),
        'test_f1_macro': float(f1_macro),
        'n_classes': int(n_classes),
        'original_feature_dim': int(X_train.shape[1]),
        'final_feature_dim': int(feature_dim),
        'ensemble': True,
        'n_ensemble_models': len(ensemble_models) if 'ensemble_models' in locals() else 1
    }
    import json
    info_path = os.path.join(output_dir, 'logistic_regression_model_info.json')
    with open(info_path, 'w') as f:
        json.dump(model_info, f, indent=2)
    print(f"✅ Model information saved: {info_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, 'logistic_regression_patient_level_results.csv')
        results_df.to_csv(results_path, index=False)
        print(f"✅ Patient-level results saved: {results_path}")

    print("\n" + "="*60)
    print("🎉 PROCESS COMPLETED!")
    print("="*60)
    print(f"\n📊 Summary:")
    print(f"  KFold CV Score: {best_cv_score:.6f}")
    print(f"  Test Accuracy: {acc:.4f} ({acc*100:.2f}%)")
    print(f"  Test F1-Score (weighted): {f1:.4f}")
    print(f"  Test F1-Score (macro): {f1_macro:.4f}")
    print(f"\nThe model is now ready for predictions.")

# Function to predict on new samples
def predict_gleason_grade(embedding_vector, 
                         model_path=os.path.join('evaluation', 'logistic_regression_results', 'logistic_regression_model.joblib'),
                         scaler_path=os.path.join('evaluation', 'logistic_regression_results', 'logistic_regression_scaler.joblib'),
                         encoder_path=os.path.join('evaluation', 'logistic_regression_results', 'logistic_regression_label_encoder.joblib'),
                         selector_path=os.path.join('evaluation', 'logistic_regression_results', 'logistic_regression_selector.joblib'),
                         pca_path=os.path.join('evaluation', 'logistic_regression_results', 'logistic_regression_pca.joblib')):
    """Predict Gleason grade for a new DINO embedding vector using Logistic Regression"""
    lr = joblib.load(model_path)
    scaler = joblib.load(scaler_path)
    label_encoder = joblib.load(encoder_path)
    selector = joblib.load(selector_path)
    
    # Reshape and scale the input
    embedding_vector = np.array(embedding_vector).reshape(1, -1)
    embedding_vector_scaled = scaler.transform(embedding_vector)
    
    # Feature selection
    embedding_vector_selected = selector.transform(embedding_vector_scaled)
    
    # PCA (varsa)
    if os.path.exists(pca_path):
        pca = joblib.load(pca_path)
        embedding_vector_final = pca.transform(embedding_vector_selected)
    else:
        embedding_vector_final = embedding_vector_selected
    
    # Get prediction and probabilities
    prediction = lr.predict(embedding_vector_final)
    probabilities = lr.predict_proba(embedding_vector_final)
    
    # 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()