import numpy as np
import pandas as pd
import sklearn
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import os
import re
import joblib
from sklearn.preprocessing import LabelEncoder
import gc  # Garbage collector
from sklearn.utils.class_weight import compute_class_weight

# 1. Load the embeddings and image paths
print("Loading embeddings and image paths...")
features = np.load('../extracted-features-test/features.npy')
with open('../extracted-features-test/image_paths.txt', 'r') as f:
    image_paths = [line.strip() for line in f.readlines()]

print(f"Loaded {len(features)} embedding vectors with shape {features.shape}")

# 2. 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)

# 3. Load the CSV data with patient information
csv_path = 'E:\\buse_thesis_prostate\\algorithms\\dinov1\\dino-code\\tcga\\prostate_tcga_wsi_paths_aws_DX_only.csv'  # Update the CSV file path
df = pd.read_csv(csv_path)
print(f"CSV loaded with {len(df)} rows")

# 4. Match patch images with case IDs and define the label per case
print("Creating case-to-grade mapping...")
case_to_grade = {}
for idx, row in df.iterrows():
    filename = row['filename']
    grade = row['gleason_grade']
    
    # Extract TCGA case ID from the CSV 'filename'
    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")

# 5. Match each embedding vector with its corresponding grade
print("Matching embeddings with grades...")
matched_features = []
matched_labels = []
matched_cases = []

# Process in batches 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]
        
        # First, try matching using the full case ID
        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 matching using the short version of the case ID (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)
    
    # Append the batch to the main lists
    matched_features.extend(batch_features)
    matched_labels.extend(batch_labels)
    matched_cases.extend(batch_cases)
    
    # Clear batch memory
    del batch_features, batch_labels, batch_cases
    gc.collect()

# 6. Reduce data size: subsample a limited number of examples per class
print(f"Total matched samples before subsampling: {len(matched_features)}")

# Convert labels to numeric IDs using LabelEncoder (to simplify processing)
label_encoder = LabelEncoder()
numeric_labels = label_encoder.fit_transform(matched_labels)
unique_labels = label_encoder.classes_

# Maximum number of samples per class
max_samples_per_class = 20000  # Adjust this value depending on your available memory

# Perform subsampling
sampled_indices = []
for label in np.unique(numeric_labels):
    # Find indices of all examples belonging to this class
    class_indices = np.where(numeric_labels == label)[0]
    
    # Random sampling (if the class has fewer than max_samples_per_class examples, take them all)
    if len(class_indices) > max_samples_per_class:
        sampled_idx = np.random.choice(class_indices, max_samples_per_class, replace=False)
    else:
        sampled_idx = class_indices
    
    sampled_indices.extend(sampled_idx)

# Create subsampled datasets
X = np.array([matched_features[i] for i in sampled_indices])
y = np.array([matched_labels[i] for i in sampled_indices])
sampled_cases = [matched_cases[i] for i in sampled_indices]

print(f"Subsampled to {len(X)} total examples")
print(f"Unique Gleason grades: {np.unique(y)}")
print(f"Class distribution after subsampling:")
for grade in np.unique(y):
    print(f"  Grade {grade}: {np.sum(y == grade)} samples")

# Clear memory
del matched_features, matched_labels, matched_cases
gc.collect()

# 7. Split into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y)
# ignoring the specific class
mask = y_train != '2+4'
X_train = X_train[mask]
y_train = y_train[mask]
mask_test = y_test != '2+4'
X_test = X_test[mask_test]
y_test = y_test[mask_test]

print(f"Training set: {X_train.shape[0]} samples")
print(f"Test set: {X_test.shape[0]} samples")

# Calculate class weights
class_weights = compute_class_weight('balanced', classes=np.unique(y_train), y=y_train)
class_weight_dict = dict(zip(np.unique(y_train), class_weights))

print("Class weights:", class_weight_dict)

# 8. Train a simple KNN model (instead of GridSearchCV)
print("Training KNN model...")
# Try several k values one by one (without GridSearchCV)
k_values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
best_k = 5  # Default value
best_score = 0

for k in k_values:
    # Basit cross-validation
    knn = KNeighborsClassifier(n_neighbors=k, weights='distance', metric='cosine')
    scores = cross_val_score(knn, X_train, y_train, cv=3, scoring='accuracy')
    avg_score = np.mean(scores)
    print(f"k={k}, Cross-validation score: {avg_score:.4f}")
    
    if avg_score > best_score:
        best_score = avg_score
        best_k = k

print(f"Best k value: {best_k} with score: {best_score:.4f}")

# Train the final model
final_model = KNeighborsClassifier(
    n_neighbors=best_k,
    weights='distance',
    metric='cosine'
)
final_model.fit(X_train, y_train)

# 9. Evaluate on test set
y_pred = final_model.predict(X_test)
print("\nClassification Report:")
print(classification_report(y_test, y_pred))

# 10. Create confusion matrix
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=np.unique(y_test),  # Change here
            yticklabels=np.unique(y_test))  # And here
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix')
plt.savefig('confusion_matrix.png')
print("Saved confusion matrix to confusion_matrix.png")
print("ROC AUC", sklearn.metrics.roc_auc_score)
# 11. Save the trained model
joblib.dump(final_model, 'tcga_gleason_knn_model.joblib')
print("Model saved to tcga_gleason_knn_model.joblib")

# 12. Save embedding-to-class mapping for future use
result_df = pd.DataFrame({
    'case_id': sampled_cases,
    'gleason_grade': y
})
result_df.to_csv('matched_cases.csv', index=False)
print("Saved case ID to class mapping in matched_cases.csv")
# Save arrays
np.save('X_train.npy', X_train)
np.save('y_train.npy', y_train)

# Save the LabelEncoder as well
joblib.dump(label_encoder, 'label_encoder.pkl')

# 13. Function to predict on new samples
def predict_gleason_grade(embedding_vector, model_path='tcga_gleason_knn_model.joblib'):
    """Predict Gleason grade for a new DINO embedding vector"""
    model = joblib.load(model_path)
    # Reshape to ensure 2D array
    embedding_vector = np.array(embedding_vector).reshape(1, -1)
    prediction = model.predict(embedding_vector)
    probabilities = model.predict_proba(embedding_vector)
    
    return {
        'predicted_grade': prediction[0],
        'probabilities': dict(zip(model.classes_, probabilities[0]))
    }

print("\nDone! You can now use the trained model for predictions.")

# Assuming y is your array of class labels
unique_classes, counts = np.unique(y, return_counts=True)

# Create a dictionary to map class labels to their counts
class_counts = dict(zip(unique_classes, counts))

# Print the counts for each class
print("Total number of samples per class:")
for class_label, count in class_counts.items():
    print(f"Class {class_label}: {count} samples")

# Print the total number of samples in the dataset
total_samples = len(y)
print(f"Total number of samples in the dataset: {total_samples}")