import numpy as np
import pandas as pd
from collections import Counter
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, classification_report, confusion_matrix
import logging
import os
from datetime import datetime
import argparse
import joblib
import json
from sklearn.impute import SimpleImputer, KNNImputer
import zarr
import wandb
import time
import psutil  # For hardware details

# 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("--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
save_data_path = args.data_path

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

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

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)

    return features, 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
if imputation_method in ['mean', 'median', 'most_frequent']:
    imputer = SimpleImputer(strategy=imputation_method)
else:
    imputer = KNNImputer(n_neighbors=5)

X_train = imputer.fit_transform(X_train)
X_test = imputer.transform(X_test)
logging.info(f"Imputed missing values using {imputation_method} strategy.")

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

# 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("Class-to-index mapping saved successfully.")

y_train = np.array([class_to_index[cls] for cls in y_train])
y_test = np.array([class_to_index[cls] for cls in y_test])

# Check Data Balance
class_counts = Counter(y_train)
class_counts_df = pd.DataFrame.from_dict(class_counts, orient='index', columns=['count'])
logging.info("Class distribution in training set:")
logging.info(class_counts_df.sort_index())

# Start tracking the training time
start_time = time.time()

# Function to train and evaluate traditional ML models
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(y_test, y_pred)
    f1 = f1_score(y_test, y_pred, average='weighted')
    precision = precision_score(y_test, y_pred, average='weighted')
    recall = recall_score(y_test, y_pred, average='weighted')
    conf_matrix = confusion_matrix(y_test, y_pred)

    # Log the performance metrics to W&B
    wandb.log({
        'accuracy': accuracy,
        'f1_score': f1,
        'precision': precision,
        'recall': recall,
        'confusion_matrix': conf_matrix.tolist()
    })
    
    # Log the performance metrics
    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}')
    logging.info(f'Confusion Matrix for {model_name}:\n{conf_matrix}')

    # Log detailed classification report
    class_report = classification_report(y_test, y_pred, target_names=[f'Class {i}' for i in range(len(unique_classes))])
    logging.info(f'Classification Report for {model_name}:\n{class_report}')
    
    # Log feature importance if the model supports it
    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}')
        
        # Optionally log the feature importances to W&B
        wandb.log({
            'feature_importances': wandb.Table(dataframe=importance_df)
        })

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

# Initialize and train the model based on the chosen architecture
if model_type == "random_forest":
    model = RandomForestClassifier(n_estimators=200, random_state=42)
    train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "RandomForest")
elif model_type == "gbm":
    model = GradientBoostingClassifier(n_estimators=100, random_state=42)
    train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "GradientBoostingMachine")
elif model_type == "logistic_regression":
    model = LogisticRegression(max_iter=2000, solver='saga', random_state=42)
    train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "LogisticRegression")
else:
    raise ValueError("Invalid model type specified. Choose from 'random_forest', 'gbm', 'logistic_regression'.")

# 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': False  # Update with GPU info if applicable
}

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

# Create metadata dictionary for 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': 200 if model_type == "random_forest" else 100 if model_type == "gbm" else None,
    'max_iter': 2000 if model_type == "logistic_regression" else None,
    'solver': 'saga' if model_type == "logistic_regression" else None,
    'timestamp': timestamp,
    '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': {str(k): convert_numpy_types(v) for k, v in dict(Counter(y_train)).items()},
    'test_class_distribution': {str(k): convert_numpy_types(v) for k, v in dict(Counter(y_test)).items()},
    'random_seed': 42,
    'training_duration': training_duration,
    'model_hyperparameters': {str(k): convert_numpy_types(v) for k, v in model.get_params().items()},
    'hardware_info': {str(k): convert_numpy_types(v) for k, v in hardware_info.items()},
}

# 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, indent=4)

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

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