import numpy as np
import pandas as pd
from collections import Counter
import xgboost as xgb  
from sklearn.linear_model import LogisticRegression
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 xgboost as xgb
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# 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=['xgboost', '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

# 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):
    """
    Load features and lccs_class data for the specified years.

    :param years: list of int, Years to load data for
    :param region: str, Region code (ES, EU, UK, AU, GLOBE)
    :param data_path: str, Base path where data is stored
    :return: tuple of np.array, Concatenated features and lccs_class for the specified years
    """
    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)}  # 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("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())

# Compute class weights
class_weight = compute_class_weight(class_weight='balanced', classes=np.unique(y_train), y=y_train)
class_weight_dict = {cls: weight for cls, weight in zip(np.unique(y_train), class_weight)}
# Compute sample weights
sample_weight = np.array([class_weight_dict[cls] for cls in y_train])

logging.info("Computed class weights for balancing.")

# Function to train and evaluate ML models with performance metrics, feature importance, and model saving
def train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, model_name):
    model.fit(X_train, y_train, sample_weight=sample_weight)
    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', zero_division=0)
    recall = recall_score(y_test, y_pred, average='weighted')

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

    # Log detailed classification report
    class_report = classification_report(y_test, 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}')

    # Log feature importance for XGBoost
    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}')

    # Log coefficients for Logistic Regression
    elif model_name == "LogisticRegression" and hasattr(model, 'coef_'):
        feature_importances = model.coef_[0]  # Assuming binary classification; adapt for multiclass
        importance_df = pd.DataFrame({
            'Feature': np.arange(len(feature_importances)),
            'Importance': feature_importances
        }).sort_values(by='Importance', ascending=False)

        logging.info(f'Feature Coefficients for {model_name}:\n{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 == "xgboost":
    model = xgb.XGBClassifier(
        n_estimators=2000,           # Number of trees
        max_depth=10,                # Maximum depth of trees
        learning_rate=0.04,          # Learning rate
        subsample=0.8,              # Subsample ratio of the training instances
        reg_alpha=0.5,               # L1 regularization
        reg_lambda=1.0,              # L2 regularization
        colsample_bytree=0.8,       # Subsample ratio of columns when constructing each tree
        eval_metric='mlogloss',     # Evaluation metric
        tree_method='hist',         # Use CPU or default tree method
        device='cuda',              # Enable GPU training
        random_state=42             # Random seed for reproducibility
    )

    train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "XGBoost")
elif model_type == "logistic_regression":
    model = LogisticRegression(max_iter=2000, solver='saga', random_state=42, class_weight='balanced')
    train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "LogisticRegression")
else:
    raise ValueError("Invalid model type specified. Choose from 'xgboost', 'logistic_regression'.")

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