import numpy as np
import pandas as pd
from collections import Counter
import xgboost as xgb  
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
from sklearn.impute import SimpleImputer, KNNImputer  
from sklearn.utils.class_weight import compute_class_weight
import zarr
import sqlite3
import sys
from codecarbon import track_emissions
import gc
import cupy as cp
# from dask_cuda import LocalCUDACluster
# from dask.distributed import Client

# # Initialize Dask CUDA cluster and client for multi-GPU usage
# cluster = LocalCUDACluster()
# client = Client(cluster)

# TODO 
# bypass for fully trained model 
# including dask training for the xgboost 
#  mutually exclusive --restart vs --resume / if restart delete the exisiting 

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../source')))
from database import log_model_run
from data_processing import apply_land_sea_mask

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a model using land use data.")
parser.add_argument("--data_path", type=str, required=True, help="Path to the directory where data is stored")
parser.add_argument("--db_path", type=str, required=True, help="Path to the SQLite database for logging model runs")
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("--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")
parser.add_argument("--model_type", type=str, required=True, choices=['xgboost', 'logistic_regression'], help="Model type to use")
parser.add_argument("--restart_training", action="store_true", help="If set, delete any existing model file and restart training from scratch")
parser.add_argument("--batch_size", type=int, default=5000, help="Batch size for processing the data")

# Parse the restart_training flag
args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
region = args.region
imputation_method = args.imputation
train_years = args.train_years
test_years = args.test_years
save_data_path = args.data_path
db_path = args.db_path
model_type = args.model_type
restart_training = args.restart_training
batch_size = args.batch_size

# 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 XGBoost model")

# Set paths for saving/loading processed data and metadata
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}")

# Compute class weights from batches rather than train_years
def compute_weights_from_generator(generator):
    y_train_all = []
    for _, y_batch in generator:
        y_train_all.extend(y_batch)
    
    y_train_all = np.array(y_train_all)
    class_weight = compute_class_weight(class_weight='balanced', classes=np.unique(y_train_all), y=y_train_all)
    class_weight_dict = {cls: weight for cls, weight in zip(np.unique(y_train_all), class_weight)}
    
    # Compute sample weights
    sample_weight = np.array([class_weight_dict[cls] for cls in y_train_all])
    
    logging.info("Computed class weights for balancing.")
    
    return class_weight_dict, sample_weight

#@track_emissions()
def train_evaluate_ml_model(model, X_test, y_test, class_to_index, model_name, 
                            save_data_path, start_year, end_year, region, timestamp):
    y_pred = model.predict(X_test)

    # Apply the water-land mask to predictions (if applicable)
    y_pred_adjusted = apply_land_sea_mask(y_pred, X_test, class_to_index)

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

    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}')

    unique_classes_in_test = np.unique(np.concatenate([y_test, y_pred_adjusted]))
    target_names = [f'Class {i}' for i in unique_classes_in_test]

    class_report = classification_report(y_test, y_pred_adjusted, target_names=target_names, zero_division=0)
    logging.info(f'Classification Report for {model_name}:\n{class_report}')

    feature_importances = None
    if hasattr(model, "feature_importances_"):
        feature_importances = model.feature_importances_
        importance_df = pd.DataFrame({
            'Feature': np.arange(len(feature_importances)),
            'Importance': feature_importances
        }).sort_values(by='Importance', ascending=False)

        logging.info(f'Feature Importances for {model_name}:\n{importance_df}')

    model_filename = f'{model_name}_{start_year}_{end_year}_{region}_{timestamp}.joblib'
    model_path = os.path.join(save_data_path, model_filename)
    joblib.dump(model, model_path)
    logging.info(f'{model_name} model saved successfully to {model_path}.')
    
    return accuracy, f1, precision, recall, feature_importances

# Adjust the data generator for handling labels
#@track_emissions()
def data_generator_with_label_adjustment(data_path, years, region, batch_size, class_to_index):
    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 = zarr.load(features_year_path)
        labels = zarr.load(lccs_class_year_path)
        
        # Log the total rows in the dataset for this year
        logging.info(f'Year: {year}, Total rows: {features.shape[0]}')

        # Now yield the data in batches
        for i in range(0, features.shape[0], batch_size):
            adjusted_labels = np.array([class_to_index[float(cls)] for cls in labels[i:i+batch_size]])
            yield features[i:i+batch_size], adjusted_labels


# Adjust the data generator without label adjustment (for computing class_to_index)
def data_generator_no_label_adjustment(data_path, years, region, batch_size):
    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 = zarr.load(features_year_path)
        labels = zarr.load(lccs_class_year_path)
        
        # Log the total rows in the dataset for this year
        logging.info(f'Year: {year}, Total rows: {features.shape[0]}')

        # Now yield the data in batches
        for i in range(0, features.shape[0], batch_size):
            yield features[i:i+batch_size], labels[i:i+batch_size]

def save_model_with_metadata(model, model_name, save_data_path, start_year, end_year, region, timestamp, feature_importances=None):
    model_filename = f'{model_name}_{start_year}_{end_year}_{region}_{timestamp}.joblib'
    model_path = os.path.join(save_data_path, model_filename)

    # Save the trained model
    joblib.dump(model, model_path)
    logging.info(f'{model_name} model saved successfully to {model_path}.')

    # Save feature importances if available
    if feature_importances is not None:
        importance_df = pd.DataFrame({
            'Feature': np.arange(len(feature_importances)),
            'Importance': feature_importances
        }).sort_values(by='Importance', ascending=False)

        importance_file = os.path.join(save_data_path, f'feature_importances_{model_name}_{timestamp}.csv')
        importance_df.to_csv(importance_file, index=False)
        logging.info(f'Feature importances saved to {importance_file}.')

    # Optionally, save any additional metadata needed
    metadata = {
        'model_name': model_name,
        'model_path': model_path,
        'timestamp': timestamp,
        'feature_importances': feature_importances if feature_importances is not None else []
    }

    metadata_file = os.path.join(save_data_path, f'model_metadata_{model_name}_{timestamp}.json')
    with open(metadata_file, 'w') as f:
        json.dump(metadata, f)
    logging.info(f'Model metadata saved to {metadata_file}.')

    return model_path, metadata_file

#@track_emissions()            
def train_xgboost_out_of_core(generator, params, num_boost_round, save_data_path, start_year, end_year, region, timestamp, batch_size=batch_size):

    model_output = os.path.join(save_data_path, f'xgboost_{start_year}_{end_year}_{region}.model')
    progress_output = os.path.join(save_data_path, f'training_progress_{start_year}_{end_year}_{region}.json')

    # Check if model and progress files exist
    model_exists = os.path.exists(model_output)
    progress_exists = os.path.exists(progress_output)

    # Load training progress (last processed batch)
    last_batch_idx = 0
    if progress_exists:
        with open(progress_output, 'r') as f:
            progress_data = json.load(f)
            last_batch_idx = progress_data.get('last_batch_idx', 0)  # Default to 0 if not found
            logging.info(f"Resuming training from batch index {last_batch_idx}.")
    else:
        logging.info("Starting training from scratch.")

    # Estimate total number of batches (if possible)
    total_batches = sum(1 for _ in generator)  # This counts total number of batches
    logging.info(f"Total number of batches to process: {total_batches}")

    # Reset the generator to start from the beginning again
    generator = data_generator_with_label_adjustment(save_data_path, train_years, region, batch_size=batch_size, class_to_index=class_to_index)

    # Training loop
    for batch_idx, (X_batch, y_batch) in enumerate(generator):
        if batch_idx < last_batch_idx:
            # Skip already processed batches
            logging.info(f"Skipping batch {batch_idx} as it has already been processed.")
            continue
        
        dtrain = xgb.DMatrix(X_batch, label=y_batch)

        if model_exists and batch_idx == last_batch_idx:
            # Resume from the existing model for the first relevant batch
            logging.info(f"Resuming training from model {model_output} at batch {batch_idx}.")
            model = xgb.train(params, dtrain, num_boost_round, xgb_model=model_output)
        else:
            # Start training from scratch or continue
            if not model_exists:
                logging.info("Starting training from scratch.")
            model = xgb.train(params, dtrain, num_boost_round)
            model.save_model(model_output)  # Save the model after each batch
            model_exists = True
        
        # Log progress after each batch
        logging.info(f"Processed batch {batch_idx + 1}/{total_batches} ({(batch_idx + 1) / total_batches * 100:.2f}% completed)")

        # Save progress after each batch
        with open(progress_output, 'w') as f:
            json.dump({'last_batch_idx': batch_idx}, f)
            logging.info(f"Saved training progress after batch {batch_idx}.")

    # Ensure the model is saved after the final batch
    model.save_model(model_output)
    model_name = 'XGBoost'
    save_model_with_metadata(model, model_name, save_data_path, start_year, end_year, region, timestamp, feature_importances=None)
    logging.info(f"Model saved to {model_output}")


# Imputation of missing values
if imputation_method in ['mean', 'median', 'most_frequent']:
    imputer = SimpleImputer(strategy=imputation_method)
else:
    imputer = KNNImputer(n_neighbors=5)

# Generate class_to_index and save it
def adjust_labels_and_compute_class_to_index(generator):
    all_classes = set()
    
    # Collect unique classes from all batches
    for _, y_batch in generator:
        unique_classes = sorted(np.unique(y_batch))
        all_classes.update(unique_classes)
    
    # Create the class_to_index mapping
    class_to_index = {float(cls): idx for idx, cls in enumerate(sorted(all_classes))}
    
    return class_to_index

# Save class_to_index
def save_class_to_index(class_to_index, save_data_path, start_year, end_year, region, timestamp):
    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("Class-to-index mapping saved successfully.")


# Main workflow
# First, compute the class_to_index mapping without adjusting labels
data_gen_for_class_mapping = data_generator_no_label_adjustment(save_data_path, train_years, region, batch_size)
class_to_index = adjust_labels_and_compute_class_to_index(data_gen_for_class_mapping)
save_class_to_index(class_to_index, save_data_path, start_year, end_year, region, timestamp)

# Re-generate the data for class weights
data_gen_for_class_weights = data_generator_with_label_adjustment(save_data_path, train_years, region, batch_size, class_to_index=class_to_index)
class_weight_dict, sample_weight = compute_weights_from_generator(data_gen_for_class_weights)

# Now re-generate the data for actual training since the generator would have been exhausted
data_gen_for_training = data_generator_with_label_adjustment(save_data_path, train_years, region, batch_size, class_to_index=class_to_index)

# Initialize the XGBoost model
params = {
    'objective': 'multi:softprob',
    'num_class': len(class_to_index),
    'max_depth': 10,
    'learning_rate': 0.04,
    'subsample': 0.8,
    'reg_alpha': 0.5,
    'reg_lambda': 1.0,
    'colsample_bytree': 0.8,
    'eval_metric': 'mlogloss',
    'tree_method': 'hist',        # Use 'hist' for tree method
    'device': 'cuda',             # Use 'cuda' for GPU
    'random_state': 42
}

# Train the model using the newly generated data
train_xgboost_out_of_core(data_gen_for_training, params, num_boost_round=2000, save_data_path=save_data_path, 
                          start_year=start_year, end_year=end_year, region=region, timestamp=timestamp)

# Evaluate the model (you can load test data fully for evaluation if needed)
X_test, y_test = next(data_generator_with_label_adjustment(save_data_path, test_years, region, batch_size, class_to_index=class_to_index))  # Load one batch for testing

# Initialize a new model for evaluation
model_path = os.path.join(save_data_path, f'xgboost_{start_year}_{end_year}_{region}.model')
model = xgb.XGBClassifier()
model.load_model(model_path)

# Use the `train_evaluate_ml_model` function for evaluation
accuracy, f1, precision, recall, feature_importances = train_evaluate_ml_model(
    model=model, 
    X_test=X_test, 
    y_test=y_test, 
    class_to_index=class_to_index, 
    model_name='XGBoost', 
    save_data_path=save_data_path, 
    start_year=start_year, 
    end_year=end_year, 
    region=region, 
    timestamp=timestamp
)

# Log the model details
model_file_path = os.path.join(save_data_path, f'xgboost_{start_year}_{end_year}_{region}_{timestamp}.joblib')

log_model_run(db_path=args.db_path, 
              model_type='xgboost', 
              data_path=save_data_path, 
              region=region, 
              start_year=start_year, 
              end_year=end_year, 
              train_years=train_years, 
              test_years=test_years, 
              timestamp=timestamp, 
              accuracy=accuracy, 
              f1=f1, 
              precision=precision, 
              recall=recall, 
              feature_importance=feature_importances,  
              class_mapping=class_to_index, 
              model_file_path=model_file_path)

logging.info("XGBoost training and evaluation complete.")