import numpy as np
import polars as pl
import pandas as pd
import torch
import gc

from torch.utils.data import DataLoader
from sklearn.preprocessing import LabelEncoder
from tqdm import tqdm
from typing import Dict

class MemoryOptimizer:        
    @staticmethod
    def cleanup_memory():
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    @staticmethod
    def get_memory_usage():
        if torch.cuda.is_available():
            return {
                'gpu_allocated': torch.cuda.memory_allocated() / 1024**3,
                'gpu_reserved': torch.cuda.memory_reserved() / 1024**3,
                'gpu_max_allocated': torch.cuda.max_memory_allocated() / 1024**3
            }
        return {'gpu_allocated': 0, 'gpu_reserved': 0, 'gpu_max_allocated': 0}
    
    @staticmethod
    def optimize_dataloader_memory(dataset, batch_size, num_workers=0):
        return DataLoader(
            dataset, 
            batch_size=batch_size, 
            shuffle=True,
            num_workers=num_workers,
            pin_memory=torch.cuda.is_available(),
            persistent_workers=False
        )
    
class UncertaintyEstimator:
    @staticmethod
    def entropy_uncertainty(probabilities: np.ndarray) -> np.ndarray:
        epsilon = 1e-8
        probabilities = np.clip(probabilities, epsilon, 1 - epsilon)
        entropy_vals = -np.sum(probabilities * np.log(probabilities), axis=1)
        return entropy_vals
    
    @staticmethod
    def variance_uncertainty(ensemble_probs: np.ndarray) -> np.ndarray:
        return np.mean(np.var(ensemble_probs, axis=0), axis=1)
    
    @staticmethod
    def mutual_information(ensemble_probs: np.ndarray) -> np.ndarray:
        mean_probs = np.mean(ensemble_probs, axis=0)
        mean_entropy = UncertaintyEstimator.entropy_uncertainty(mean_probs)
        
        individual_entropies = np.array([
            UncertaintyEstimator.entropy_uncertainty(probs) 
            for probs in ensemble_probs
        ])
        expected_entropy = np.mean(individual_entropies, axis=0)
        
        return mean_entropy - expected_entropy
    
    def estimate_uncertainty(self, ensemble_predictions: np.ndarray, 
                           ensemble_probabilities: np.ndarray) -> Dict[str, np.ndarray]:
        mean_probs = np.mean(ensemble_probabilities, axis=0)
        
        uncertainties = {
            'entropy': self.entropy_uncertainty(mean_probs),
            'variance': self.variance_uncertainty(ensemble_probabilities),
            'mutual_info': self.mutual_information(ensemble_probabilities),
        }
        
        uncertainties['total'] = uncertainties['entropy'] + uncertainties['variance']
        
        max_entropy = np.log(mean_probs.shape[1])  # log(num_classes)
        uncertainties['confidence'] = 1 - (uncertainties['entropy'] / max_entropy)
        
        return uncertainties

def preprocess_ultra_fast(test_df, target_col='class'):
    X_test = test_df.drop(columns=[target_col], errors='ignore')
    y_test = test_df[target_col] if target_col in test_df.columns else None
    
    print("Encoding target labels...")
    label_encoder = LabelEncoder()
    
    if y_test is not None:
        # Fit the encoder first, then transform
        label_encoder.fit(y_test)
        y_test_encoded = label_encoder.transform(y_test)
    else:
        y_test_encoded = None
        
    
    # Identify categorical and numerical columns
    categorical_cols = X_test.select_dtypes(include=['object', 'category']).columns.tolist()
    numerical_cols = X_test.select_dtypes(include=[np.number]).columns.tolist()
    
    print(f"Categorical columns ({len(categorical_cols)}): {categorical_cols}")
    print(f"Numerical columns ({len(numerical_cols)}): Found {len(numerical_cols)} numerical columns")
    
    # SUPER FAST CATEGORICAL ENCODING - Only for categorical columns
    print("Fast categorical encoding...")
    cat_encoders = {}
    
    if categorical_cols:
        with tqdm(total=len(categorical_cols), desc="Encoding categorical") as pbar:
            for col in categorical_cols:
                # Get unique values efficiently
                unique_vals = set()
                if col in X_test.columns:
                    unique_vals.update(X_test[col].astype(str).unique())
                
                
                unique_list = sorted(list(unique_vals))
                mapping = {val: idx for idx, val in enumerate(unique_list)}
                
                # Apply mapping
                X_test[col] = X_test[col].astype(str).map(mapping)
                
                # Store reverse mapping for later use
                reverse_mapping = {idx: val for val, idx in mapping.items()}
                cat_encoders[col] = {'mapping': mapping, 'reverse_mapping': reverse_mapping}
                
                pbar.update(1)
    
    scaler = None

    preprocessors = {
        'feature_engineer': None,  # Skip feature engineering
        'label_encoder': label_encoder,
        'categorical_encoders': cat_encoders,
        'scaler': scaler,
        'categorical_cols': categorical_cols,
        'numerical_cols': numerical_cols
    }
    
    print(f"Final shapes - Test: {X_test.shape}")
    
    return (X_test, y_test_encoded, categorical_cols, numerical_cols, preprocessors)

def load_and_clean_data(path, columns_to_drop, drugname_a, drugname_b):
    
    df = pd.read_csv(path)
    filtered_df = df[
        (df["drugname-drug_a"] == drugname_a) & 
        (df["drugname-drug_b"] == drugname_b)
    ]
    existing_cols_to_drop = [col for col in columns_to_drop if col in df.columns]
    if existing_cols_to_drop:
        filtered_df = filtered_df.drop(existing_cols_to_drop)
    # result = filtered_df.to_pandas()
    del df  # Free memory
    MemoryOptimizer.cleanup_memory()
    return filtered_df
    # except Exception as e:
    #     print(f"Error loading {path}: {e}")
    #     return None