import numpy as np
import pandas as pd
from collections import Counter
from sklearn.impute import SimpleImputer, KNNImputer
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score, f1_score, precision_score, recall_score
import logging
import os
from datetime import datetime
import argparse
import joblib
import json
import zarr
import cuml
from cuml.ensemble import RandomForestClassifier as cuRFClassifier
import cupy as cp
import wandb
import time  # For capturing training duration
import psutil  # For hardware details
import json
import gc

# Hyperparameters 
N_ESTIMATORS = 200
MAX_DEPTH = 1000
MIN_SAMPLES_SPLIT = 2
MIN_SAMPLES_LEAF = 1
MAX_FEATURES = 'sqrt'  # Updated from 'auto' to 'sqrt' to avoid deprecation

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a model using land use data.")
parser.add_argument("--start_year", type=int, required=True, help="Start year for the data range")
parser.add_argument("--end_year", type=int, required=True, help="End year for the data range")
parser.add_argument("--region", type=str, required=True, help="Region for the data slice (ES, EU, UK, AU, GLOBE)")
parser.add_argument("--model_type", type=str, required=True, choices=['random_forest', 'gbm', 'logistic_regression'], help="Model type to use")
parser.add_argument("--imputation", type=str, choices=['mean', 'median', 'most_frequent', 'knn'], default='mean', help="Imputation method for handling NaN values")
parser.add_argument("--train_years", nargs='+', type=int, required=True, help="Years to use for training")
parser.add_argument("--test_years", nargs='+', type=int, required=True, help="Years to use for testing")

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
region = args.region
model_type = args.model_type
imputation_method = args.imputation
train_years = args.train_years
test_years = args.test_years

# Initialize W&B
wandb.init(
    project="land-use-classification",
    config={
        "start_year": start_year,
        "end_year": end_year,
        "region": region,
        "model_type": model_type,
        "imputation": imputation_method,
        "train_years": train_years,
        "test_years": test_years,
        "n_estimators": N_ESTIMATORS,
        "max_depth": MAX_DEPTH,
        "min_samples_split": MIN_SAMPLES_SPLIT,
        "min_samples_leaf": MIN_SAMPLES_LEAF,
        "max_features": MAX_FEATURES
    },
    mode="offline"  # Run in offline mode
)

def log_gpu_memory():
    mempool = cp.get_default_memory_pool()
    used_bytes = mempool.used_bytes()
    total_bytes = mempool.total_bytes()
    print(f"Memory used: {used_bytes / (1024**2):.2f} MB, Total allocated: {total_bytes / (1024**2):.2f} MB")
    logging.info(f"Memory used: {used_bytes / (1024**2):.2f} MB, Total allocated: {total_bytes / (1024**2):.2f} MB")
    
def predict_in_batches(model, X_test, batch_size=20):  # Trying smaller batches
    n_samples = X_test.shape[0]
    y_pred = []
    
    for start in range(0, n_samples, batch_size):
        end = min(start + batch_size, n_samples)
        X_test_batch = X_test[start:end]
        
        # Log memory before batch prediction
        logging.info(f"Batch {start // batch_size + 1}: Memory usage before prediction:")
        log_gpu_memory()

        # Memory cleanup and synchronization
        cp.cuda.Stream.null.synchronize()
        cp._default_memory_pool.free_all_blocks()

        try:
            # Predict on the current batch
            y_pred_batch = model.predict(X_test_batch)
            y_pred.append(y_pred_batch)

            # Log memory after batch prediction
            logging.info(f"Batch {start // batch_size + 1}: Memory usage after prediction:")
            log_gpu_memory()

        except MemoryError:
            logging.error("Memory allocation error during batch prediction.")
            break

        # Force memory cleanup after batch
        cp._default_memory_pool.free_all_blocks()
        gc.collect()  # Explicit garbage collection for any remaining objects in memory

        # Log memory after cleanup
        logging.info(f"Batch {start // batch_size + 1}: Memory usage after cleanup:")
        log_gpu_memory()

    # Concatenate all predictions if available, else return None
    return cp.concatenate(y_pred) if y_pred else None



# Configure logging
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
log_file_name = f"train_log_{start_year}_{end_year}_{region}_{timestamp}.log"
log_file_path = os.path.join(logs_path, log_file_name)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[
    logging.FileHandler(log_file_path),
    logging.StreamHandler()
])
logging.info(f"Training model of type: {model_type}")
logging.info(f"Hyperparameters: N_ESTIMATORS={N_ESTIMATORS}, MAX_DEPTH={MAX_DEPTH}, MIN_SAMPLES_SPLIT={MIN_SAMPLES_SPLIT}, "
             f"MIN_SAMPLES_LEAF={MIN_SAMPLES_LEAF}, MAX_FEATURES={MAX_FEATURES}")

# Set paths for saving/loading processed data and metadata
save_data_path = '/gpfs/projects/bsc32/cerise/outputs/bsc032013/ml/'
metadata_path = os.path.join(save_data_path, f'metadata_{start_year}_{end_year}_{region}.json')

# Load metadata
with open(metadata_path, 'r') as f:
    metadata = json.load(f)

# Log the loaded metadata
logging.info(f"Loaded metadata: {metadata}")

# Convert your data to GPU format (using CuPy for arrays)
def to_cupy(data):
    return cp.array(data)

# Modify your load_data_for_years function to return CuPy arrays:
def load_data_for_years(years, region, data_path):
    features_list = []
    lccs_class_list = []

    for year in years:
        features_year_path = os.path.join(data_path, f'features_{year}_{region}.zarr')
        lccs_class_year_path = os.path.join(data_path, f'lccs_class_{year}_{region}.zarr')

        features_list.append(zarr.load(features_year_path))
        lccs_class_list.append(zarr.load(lccs_class_year_path))

    features = np.concatenate(features_list, axis=0)
    lccs_class = np.concatenate(lccs_class_list, axis=0)

    # Convert features and lccs_class to GPU (CuPy) arrays
    return to_cupy(features), to_cupy(lccs_class)

# Load data for training and testing years with mixed precision (float16)
X_train, y_train = load_data_for_years(train_years, region, save_data_path)
X_test, y_test = load_data_for_years(test_years, region, save_data_path)

# Impute missing values using Scikit-learn as cuML does not have imputation yet
if imputation_method in ['mean', 'median', 'most_frequent']:
    imputer = SimpleImputer(strategy=imputation_method)
else:
    imputer = KNNImputer(n_neighbors=5)

# Apply imputation and convert back to CuPy (with mixed precision float16)
X_train = to_cupy(imputer.fit_transform(cp.asnumpy(X_train)))
X_test = to_cupy(imputer.transform(cp.asnumpy(X_test)))

# Ensure the data is in float16 format to save memory
X_train = X_train.astype(cp.float16)
X_test = X_test.astype(cp.float16)


# Adjust labels to ensure they start from 0 and go up to n_classes-1
unique_classes = sorted(np.unique(y_train))
class_to_index = {float(cls): idx for idx, cls in enumerate(unique_classes)}  # Convert keys to float

# Save class_to_index to a JSON file
with open(os.path.join(save_data_path, f'class_to_index_{start_year}_{end_year}_{region}_{timestamp}.json'), 'w') as f:
    json.dump(class_to_index, f)
    logging.info(f"class_to_index.json is created in file name is:{class_to_index}")

# Capture start time for tracking training duration
start_time = time.time()

def train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, model_name, batch_size=20):
    # Train the model
    model.fit(X_train, y_train)

    # Predict in batches to avoid GPU memory issues
    y_pred = predict_in_batches(model, X_test, batch_size=batch_size)

    # Compute performance metrics
    accuracy = accuracy_score(cp.asnumpy(y_test), cp.asnumpy(y_pred))
    f1 = f1_score(cp.asnumpy(y_test), cp.asnumpy(y_pred), average='weighted')
    precision = precision_score(cp.asnumpy(y_test), cp.asnumpy(y_pred), average='weighted', zero_division=0)
    recall = recall_score(cp.asnumpy(y_test), cp.asnumpy(y_pred), average='weighted')

    # Log the performance metrics to W&B
    wandb.log({
        'accuracy': accuracy,
        'f1_score': f1,
        'precision': precision,
        'recall': recall
    })

    # Log the performance metrics in your log file
    logging.info(f'Accuracy of {model_name}: {accuracy*100:.2f}%')
    logging.info(f'F1 Score of {model_name}: {f1:.2f}')
    logging.info(f'Precision of {model_name}: {precision:.2f}')
    logging.info(f'Recall of {model_name}: {recall:.2f}')

    # Log detailed classification report
    class_report = classification_report(cp.asnumpy(y_test), cp.asnumpy(y_pred), target_names=[f'Class {i}' for i in range(len(unique_classes))], zero_division=0)
    logging.info(f'Classification Report for {model_name}:\n{class_report}')

    return accuracy, f1, precision, recall, class_report

# Initialize and train the RAPIDS RandomForestClassifier
model = cuRFClassifier(n_estimators=N_ESTIMATORS, random_state=42, max_depth=MAX_DEPTH,
                       min_samples_split=MIN_SAMPLES_SPLIT, min_samples_leaf=MIN_SAMPLES_LEAF, max_features=MAX_FEATURES)

# Train the model and collect metrics
accuracy, f1, precision, recall, class_report, conf_matrix, model_file_path = train_evaluate_ml_model(
    model, X_train, X_test, y_train, y_test, "cuML_RandomForest"
)

# Capture training duration
training_duration = time.time() - start_time
logging.info(f"Training completed in {training_duration:.2f} seconds.")

# Capture hardware details
hardware_info = {
    'num_cores': os.cpu_count(),
    'memory_gb': psutil.virtual_memory().total / (1024 ** 3),
    'gpu_used': True  # Set to True since we are using GPU for RAPIDS
}

# Function to convert numpy and cupy data types to standard Python types
def convert_numpy_types(obj):
    if isinstance(obj, (np.integer, cp.integer)):
        return int(obj)
    elif isinstance(obj, (np.floating, cp.floating)):
        return float(obj)
    elif isinstance(obj, (np.ndarray, cp.ndarray)):
        return obj.tolist()  # Convert numpy and cupy arrays to lists
    else:
        return obj

# Convert cupy array to numpy before using Counter
train_class_distribution = {str(k): convert_numpy_types(v) for k, v in dict(Counter(cp.asnumpy(y_train))).items()}
test_class_distribution = {str(k): convert_numpy_types(v) for k, v in dict(Counter(cp.asnumpy(y_test))).items()}

# Function to filter out non-serializable types
def serialize_params(params):
    serializable_params = {}
    for k, v in params.items():
        try:
            json.dumps(v)  # Test if the value is JSON serializable
            serializable_params[k] = v
        except TypeError:
            serializable_params[k] = str(v)  # Convert non-serializable values to string
    return serializable_params

# Create metadata dictionary for the training process
training_metadata = {
    'start_year': start_year,
    'end_year': end_year,
    'region': region,
    'model_type': model_type,
    'train_years': [convert_numpy_types(year) for year in train_years],
    'test_years': [convert_numpy_types(year) for year in test_years],
    'imputation': imputation_method,
    'n_estimators': N_ESTIMATORS,
    'max_depth': MAX_DEPTH,
    'min_samples_split': MIN_SAMPLES_SPLIT,
    'min_samples_leaf': MIN_SAMPLES_LEAF,
    'class_to_index': {str(k): v for k, v in class_to_index.items()},
    'X_train_shape': convert_numpy_types(X_train.shape),
    'X_test_shape': convert_numpy_types(X_test.shape),
    'train_class_distribution': train_class_distribution,
    'test_class_distribution': test_class_distribution,
    'random_seed': 42,
    'training_duration': training_duration,
    'model_hyperparameters': serialize_params(model.get_params()),  # Ensure hyperparameters are serializable
    'accuracy': accuracy,
    'f1_score': f1,
    'precision': precision,
    'recall': recall,
    'classification_report': class_report,
    'confusion_matrix': conf_matrix.tolist(),  # Ensure it's serializable
    'hardware_info': {str(k): convert_numpy_types(v) for k, v in hardware_info.items()},
    'model_file_path': model_file_path
}

# Save the metadata to a JSON file
training_metadata_filename = f'{model_type}_training_metadata_{start_year}_{end_year}_{region}_{timestamp}.json'
training_metadata_path = os.path.join(save_data_path, training_metadata_filename)

with open(training_metadata_path, 'w') as f:
    json.dump(training_metadata, f, indent=4)

logging.info(f'Training metadata saved successfully to {training_metadata_path}')
wandb.finish()
