import numpy as np
import pandas as pd
from collections import Counter
from sklearn.impute import SimpleImputer, KNNImputer  # Import Imputers
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, classification_report
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

# Hyperparameters
N_ESTIMATORS = 500
MAX_DEPTH = 50
MIN_SAMPLES_SPLIT = 20
MIN_SAMPLES_LEAF = 10
MAX_FEATURES = 'log2' #sqrt log2 int-value  float-value 

# 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
#  to sync, please use the login node and run the following commond 
#   wandb sync wandb/offline-run-*
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
)

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

# Note: Convert the data back to NumPy before imputing as Scikit-learn uses CPU
X_train = imputer.fit_transform(cp.asnumpy(X_train))
X_test = imputer.transform(cp.asnumpy(X_test))

# Convert back to CuPy arrays after imputation
X_train = to_cupy(X_train)
X_test = to_cupy(X_test)

# Ensure the data is in float32 format for cuML and pickling compatibility
X_train = X_train.astype(cp.float32)
X_test = X_test.astype(cp.float32)

# 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}")

def train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, model_name):
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)

    # 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)  # Handle division by zero
    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 with zero_division handling
    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}')

    # Save the model
    model_name = f'{model_name}_{start_year}_{end_year}_{region}_{timestamp}.joblib'
    joblib.dump(model, os.path.join(save_data_path, model_name))
    logging.info(f'{model_name} model saved successfully.')

    # Save the model to W&B
    wandb.save(os.path.join(save_data_path, model_name))


# Initialize and train the RAPIDS RandomForestClassifier
if model_type == "random_forest":
    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_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "cuML_RandomForest")

# Create metadata dictionary for training process
training_metadata = {
    'start_year': start_year,
    'end_year': end_year,
    'region': region,
    'model_type': model_type,
    'train_years': train_years,
    'test_years': 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,
    'max_features': MAX_FEATURES,
    'timestamp': timestamp,
    'class_to_index': class_to_index,  # Save class mapping information
}

# Save training 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)
    logging.info(f'Training metadata saved to {training_metadata_path}')

logging.info("Training and evaluation complete.")
wandb.finish()
